Performance Optimisation Around Dependency Calculations

go/vt/vtgate/planbuilder/operators/ast_to_op.go

@@ -120,17 +120,18 @@ func cloneASTAndSemState[T sqlparser.SQLNode](ctx *plancontext.PlanningContext,
 }
 
 // findTablesContained returns the TableSet of all the contained
-func findTablesContained(ctx *plancontext.PlanningContext, node sqlparser.SQLNode) (result semantics.TableSet) {
+func findTablesContained(ctx *plancontext.PlanningContext, node sqlparser.SQLNode) semantics.TableSet {
+	var tables semantics.MutableTableSet
 	_ = sqlparser.Walk(func(node sqlparser.SQLNode) (bool, error) {
 		t, ok := node.(*sqlparser.AliasedTableExpr)
 		if !ok {
 			return true, nil
 		}
 		ts := ctx.SemTable.TableSetFor(t)
-		result = result.Merge(ts)
+		tables.MergeInPlace(ts)
 		return true, nil
 	}, node)
-	return
+	return tables.ToImmutable()
 }
 
 // joinPredicateCollector is used to inspect the predicates inside the subquery, looking for any

go/vt/vtgate/planbuilder/operators/helpers.go

@@ -71,14 +71,15 @@ type tableIDIntroducer interface {
 	introducesTableID() semantics.TableSet
 }
 
-func TableID(op Operator) (result semantics.TableSet) {
+func TableID(op Operator) semantics.TableSet {
+	var tables semantics.MutableTableSet
 	_ = Visit(op, func(this Operator) error {
 		if tbl, ok := this.(tableIDIntroducer); ok {
-			result = result.Merge(tbl.introducesTableID())
+			tables.MergeInPlace(tbl.introducesTableID())
 		}
 		return nil
 	})
-	return
+	return tables.ToImmutable()
 }
 
 // TableUser is used to signal that this operator directly interacts with one or more tables

go/vt/vtgate/planbuilder/operators/querygraph.go

@@ -66,11 +66,11 @@ var _ Operator = (*QueryGraph)(nil)
 
 // Introduces implements the tableIDIntroducer interface
 func (qg *QueryGraph) introducesTableID() semantics.TableSet {
-	var ts semantics.TableSet
+	var ts semantics.MutableTableSet
 	for _, table := range qg.Tables {
-		ts = ts.Merge(table.ID)
+		ts.MergeInPlace(table.ID)
 	}
-	return ts
+	return ts.ToImmutable()
 }
 
 // GetPredicates returns the predicates that are applicable for the two given TableSets

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

@@ -806,9 +806,9 @@ func (r *Route) getTruncateColumnCount() int {
 }
 
 func (r *Route) introducesTableID() semantics.TableSet {
-	id := semantics.EmptyTableSet()
+	var tables semantics.MutableTableSet
 	for _, route := range r.MergedWith {
-		id = id.Merge(TableID(route))
+		tables.MergeInPlace(TableID(route))
 	}
-	return id
+	return tables.ToImmutable()
 }

go/vt/vtgate/planbuilder/operators/subquery_planning.go

@@ -211,7 +211,7 @@ We transform it to:
         └─────┘                                                 └─────┘
 We are rewriting all expressions in the subquery to use arguments any columns
 coming from the LHS. The join predicate is not affected, but we are adding
-any new columns needed by the inner subquery to the JoinVars that the join
+any new columns needed by the inner subquery to the JoinVars that the join1
 will handle.
 */
 func tryPushSubQueryInJoin(
@@ -228,11 +228,11 @@ func tryPushSubQueryInJoin(
 	// we want to push the subquery as close to its needs
 	// as possible, so that we can potentially merge them together
 	// TODO: we need to check dependencies and break apart all expressions in the subquery, not just the merge predicates
-	deps := semantics.EmptyTableSet()
+	var tables semantics.MutableTableSet
 	for _, predicate := range inner.GetMergePredicates() {
-		deps = deps.Merge(ctx.SemTable.RecursiveDeps(predicate))
+		tables.MergeInPlace(ctx.SemTable.RecursiveDeps(predicate))
 	}
-	deps = deps.Remove(innerID)
+	deps := tables.ToImmutable().Remove(innerID)
 
 	// in general, we don't want to push down uncorrelated subqueries into the RHS of a join,
 	// since this side is executed once per row from the LHS, so we would unnecessarily execute

go/vt/vtgate/planbuilder/plan_test.go

@@ -872,3 +872,41 @@ func benchmarkPlanner(b *testing.B, version plancontext.PlannerVersion, testCase
 		}
 	}
 }
+
+func (s *planTestSuite) TestMy() {
+	vschema, err := vschemawrapper.NewVschemaWrapper(
+		vtenv.NewTestEnv(),
+		loadSchema(s.T(), "vschemas/my_schema.json", true),
+		TestBuilder,
+	)
+	require.NoError(s.T(), err)
+	// vschema := &vschemawrapper.VSchemaWrapper{
+	// 	V:             loadSchema(s.T(), "vschemas/my_schema.json", true),
+	// 	SysVarEnabled: true,
+	// 	Version:       Gen4,
+	// 	Env:           vtenv.NewTestEnv(),
+	// }
+
+	s.testFile("my.json", vschema, false)
+}
+
+func BenchmarkMine(b *testing.B) {
+	vschema, err := vschemawrapper.NewVschemaWrapper(
+		vtenv.NewTestEnv(),
+		loadSchema(b, "vschemas/my_schema.json", true),
+		TestBuilder,
+	)
+	require.NoError(b, err)
+	testCases := readJSONTests("my.json")
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for range b.N {
+		for _, tcase := range testCases {
+			plan, _ := TestBuilder(tcase.Query, vschema, vschema.CurrentDb())
+			if plan == nil {
+				panic("")
+			}
+		}
+	}
+}

go/vt/vtgate/planbuilder/testdata/vschemas/my_schema.json

@@ -0,0 +1,110 @@
+{
+  "keyspaces": {
+    "unsharded_ns": {
+      "sharded": false,
+      "tables": {
+        "my_table_seq": {
+          "type": "sequence",
+          "columns": [
+            {
+              "name": "next_id",
+              "type": "UINT64"
+            },
+            {
+              "name": "cache",
+              "type": "UINT32"
+            },
+            {
+              "name": "id",
+              "type": "UINT64"
+            }
+          ],
+          "column_list_authoritative": true
+        }
+      }
+    },
+    "main": {
+      "sharded": true,
+      "vindexes": {
+        "binary_vdx": {
+          "type": "binary"
+        }
+      },
+      "tables": {
+        "my_table": {
+          "column_vindexes": [
+            {
+              "columns": [
+                "sdKey"
+              ],
+              "name": "binary_vdx"
+            }
+          ],
+          "auto_increment": {
+            "column": "id",
+            "sequence": "my_table_seq"
+          },
+          "columns": [
+            {
+              "name": "col01",
+              "type": "VARCHAR"
+            },
+            {
+              "name": "col02",
+              "type": "INT64"
+            },
+            {
+              "name": "col03",
+              "type": "BLOB"
+            },
+            {
+              "name": "col04",
+              "type": "VARCHAR"
+            },
+            {
+              "name": "col05",
+              "type": "VARBINARY"
+            },
+            {
+              "name": "col06",
+              "type": "INT8"
+            },
+            {
+              "name": "col07",
+              "type": "INT64"
+            },
+            {
+              "name": "col08",
+              "type": "VARBINARY"
+            },
+            {
+              "name": "sdKey",
+              "type": "VARBINARY"
+            },
+            {
+              "name": "col09",
+              "type": "INT64"
+            },
+            {
+              "name": "col10",
+              "type": "INT64"
+            },
+            {
+              "name": "col11",
+              "type": "INT64"
+            },
+            {
+              "name": "col12",
+              "type": "INT64"
+            },
+            {
+              "name": "id",
+              "type": "UINT64"
+            }
+          ],
+          "column_list_authoritative": true
+        }
+      }
+    }
+  }
+}

go/vt/vtgate/semantics/analyzer.go

@@ -174,7 +174,7 @@ func (a *analyzer) newSemTable(
 		Direct:                    a.binder.direct,
 		ExprTypes:                 a.typer.m,
 		Tables:                    a.tables.Tables,
-		DMLTargets:                a.binder.targets,
+		DMLTargets:                a.binder.targets.ToImmutable(),
 		NotSingleRouteErr:         a.notSingleRouteErr,
 		NotUnshardedErr:           a.unshardedErr,
 		Warning:                   a.warning,

go/vt/vtgate/semantics/binder.go

@@ -31,7 +31,7 @@ import (
 type binder struct {
 	recursive ExprDependencies
 	direct    ExprDependencies
-	targets   TableSet
+	targets   MutableTableSet
 	scoper    *scoper
 	tc        *tableCollector
 	org       originable
@@ -81,7 +81,7 @@ func (b *binder) bindUpdateExpr(ue *sqlparser.UpdateExpr) error {
 	if !ok {
 		return nil
 	}
-	b.targets = b.targets.Merge(ts)
+	b.targets.MergeInPlace(ts)
 	return nil
 }
 
@@ -96,7 +96,7 @@ func (b *binder) bindTableNames(cursor *sqlparser.Cursor, tables sqlparser.Table
 		if err != nil {
 			return err
 		}
-		b.targets = b.targets.Merge(finalDep.direct)
+		b.targets.MergeInPlace(finalDep.direct)
 	}
 	return nil
 }
@@ -184,19 +184,20 @@ func (b *binder) findDependentTableSet(current *scope, target sqlparser.TableNam
 
 func (b *binder) bindCountStar(node *sqlparser.CountStar) error {
 	scope := b.scoper.currentScope()
-	var ts TableSet
+	var tables MutableTableSet
 	for _, tbl := range scope.tables {
 		switch tbl := tbl.(type) {
 		case *vTableInfo:
 			for _, col := range tbl.cols {
 				if sqlparser.Equals.Expr(node, col) {
-					ts = ts.Merge(b.recursive[col])
+					tables.MergeInPlace(b.recursive[col])
 				}
 			}
 		default:
-			ts = ts.Merge(tbl.getTableSet(b.org))
+			tables.MergeInPlace(tbl.getTableSet(b.org))
 		}
 	}
+	ts := tables.ToImmutable()
 	b.recursive[node] = ts
 	b.direct[node] = ts
 	return nil
@@ -246,15 +247,16 @@ func (b *binder) setSubQueryDependencies(subq *sqlparser.Subquery) error {
 	subqRecursiveDeps := b.recursive.dependencies(subq)
 	subqDirectDeps := b.direct.dependencies(subq)
 
-	tablesToKeep := EmptyTableSet()
+	var tables MutableTableSet
 	sco := currScope
 	for sco != nil {
 		for _, table := range sco.tables {
-			tablesToKeep = tablesToKeep.Merge(table.getTableSet(b.org))
+			tables.MergeInPlace(table.getTableSet(b.org))
 		}
 		sco = sco.parent
 	}
 
+	tablesToKeep := tables.ToImmutable()
 	b.recursive[subq] = subqRecursiveDeps.KeepOnly(tablesToKeep)
 	b.direct[subq] = subqDirectDeps.KeepOnly(tablesToKeep)
 	return nil

go/vt/vtgate/semantics/bitset/mutable.go

@@ -0,0 +1,46 @@
+/*
+Copyright 2024 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 bitset
+
+// Mutable is a growable, in-place version that we can OR bits into
+// and eventually turn into a stable (immutable) TableSet.
+type Mutable struct {
+	data []byte
+}
+
+// Or merges another TableSet into this Mutable, resizing if needed.
+func (m *Mutable) Or(ts Bitset) {
+	// If ts is longer than our current data, grow to accommodate it.
+	if len(ts) > len(m.data) {
+		oldData := m.data
+		m.data = make([]byte, len(ts))
+		copy(m.data, oldData)
+	}
+	// Merge in-place.
+	for i := 0; i < len(ts); i++ {
+		m.data[i] |= ts[i]
+	}
+}
+
+// AsImmutable finalizes the Mutable into a TableSet, trimming trailing zeros.
+func (m *Mutable) AsImmutable() Bitset {
+	trim := len(m.data)
+	for trim > 0 && m.data[trim-1] == 0 {
+		trim--
+	}
+	return toBitset(m.data[:trim])
+}

go/vt/vtgate/semantics/cte_table.go

@@ -163,10 +163,11 @@ type CTE struct {
 	Merged bool
 }
 
-func (cte *CTE) recursive(org originable) (id TableSet) {
+func (cte *CTE) recursive(org originable) TableSet {
 	if cte.recursiveDeps != nil {
 		return *cte.recursiveDeps
 	}
+	var id MutableTableSet
 
 	// We need to find the recursive dependencies of the CTE
 	// We'll do this by walking the inner query and finding all the tables
@@ -175,8 +176,8 @@ func (cte *CTE) recursive(org originable) (id TableSet) {
 		if !ok {
 			return true, nil
 		}
-		id = id.Merge(org.tableSetFor(ate))
+		id.MergeInPlace(org.tableSetFor(ate))
 		return true, nil
 	}, cte.Query)
-	return
+	return id.ToImmutable()
 }