wip: fewer allocations

Signed-off-by: Vicent Marti <[email protected]>

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

@@ -18,7 +18,6 @@ package operators
 
 import (
 	"fmt"
-	"vitess.io/vitess/go/vt/vtgate/semantics/bitset"
 
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
@@ -122,17 +121,17 @@ func cloneASTAndSemState[T sqlparser.SQLNode](ctx *plancontext.PlanningContext,
 
 // findTablesContained returns the TableSet of all the contained
 func findTablesContained(ctx *plancontext.PlanningContext, node sqlparser.SQLNode) semantics.TableSet {
-	bs := bitset.NewMutable()
+	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)
-		bs.Or(bitset.Bitset(ts))
+		tables.MergeInPlace(ts)
 		return true, nil
 	}, node)
-	return semantics.TableSet(bs.AsImmutable())
+	return tables.ToImmutable()
 }
 
 // joinPredicateCollector is used to inspect the predicates inside the subquery, looking for any

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

@@ -72,14 +72,14 @@ type tableIDIntroducer interface {
 }
 
 func TableID(op Operator) semantics.TableSet {
-	var tables []semantics.TableSet
+	var tables semantics.MutableTableSet
 	_ = Visit(op, func(this Operator) error {
 		if tbl, ok := this.(tableIDIntroducer); ok {
-			tables = append(tables, tbl.introducesTableID())
+			tables.MergeInPlace(tbl.introducesTableID())
 		}
 		return nil
 	})
-	return semantics.MergeTableSets(tables...)
+	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 = append(ts, table.ID)
+		ts.MergeInPlace(table.ID)
 	}
-	return semantics.MergeTableSets(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 {
-	var ts []semantics.TableSet
+	var ts semantics.MutableTableSet
 	for _, route := range r.MergedWith {
-		ts = append(ts, TableID(route))
+		ts.MergeInPlace(TableID(route))
 	}
-	return semantics.MergeTableSets(ts...)
+	return ts.ToImmutable()
 }

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

@@ -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
-	var ts []semantics.TableSet
+	var ts semantics.MutableTableSet
 	for _, predicate := range inner.GetMergePredicates() {
-		ts = append(ts, ctx.SemTable.RecursiveDeps(predicate))
+		ts.MergeInPlace(ctx.SemTable.RecursiveDeps(predicate))
 	}
-	deps := semantics.MergeTableSets(ts...).Remove(innerID)
+	deps := ts.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
@@ -377,8 +377,8 @@ func rewriteOriginalPushedToRHS(ctx *plancontext.PlanningContext, expression sql
 // this is used when we push an operator from above the subquery into the outer side of the subquery
 func rewriteColNameToArgument(
 	ctx *plancontext.PlanningContext,
-	in sqlparser.Expr, // the expression to rewrite
-	se SubQueryExpression, // the subquery expression we are rewriting
+	in sqlparser.Expr,       // the expression to rewrite
+	se SubQueryExpression,   // the subquery expression we are rewriting
 	subqueries ...*SubQuery, // the inner subquery operators
 ) sqlparser.Expr {
 	// the visitor function that will rewrite the expression tree

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
 }
 
@@ -91,15 +91,15 @@ func (b *binder) bindTableNames(cursor *sqlparser.Cursor, tables sqlparser.Table
 		return nil
 	}
 	current := b.scoper.currentScope()
-	var targets []TableSet
+	var targets MutableTableSet
 	for _, target := range tables {
 		finalDep, err := b.findDependentTableSet(current, target)
 		if err != nil {
 			return err
 		}
-		targets = append(targets, finalDep.direct)
+		targets.MergeInPlace(finalDep.direct)
 	}
-	b.targets = MergeTableSets(targets...)
+	b.targets = targets
 	return nil
 }
 
@@ -186,21 +186,21 @@ func (b *binder) findDependentTableSet(current *scope, target sqlparser.TableNam
 
 func (b *binder) bindCountStar(node *sqlparser.CountStar) error {
 	scope := b.scoper.currentScope()
-	var deps []TableSet
+	var deps MutableTableSet
 	for _, tbl := range scope.tables {
 		switch tbl := tbl.(type) {
 		case *vTableInfo:
 			for _, col := range tbl.cols {
 				if sqlparser.Equals.Expr(node, col) {
-					deps = append(deps, b.recursive[col])
+					deps.MergeInPlace(b.recursive[col])
 				}
 			}
 		default:
-			deps = append(deps, tbl.getTableSet(b.org))
+			deps.MergeInPlace(tbl.getTableSet(b.org))
 		}
 	}
 
-	ts := MergeTableSets(deps...)
+	ts := deps.ToImmutable()
 	b.recursive[node] = ts
 	b.direct[node] = ts
 	return nil
@@ -250,16 +250,16 @@ func (b *binder) setSubQueryDependencies(subq *sqlparser.Subquery) error {
 	subqRecursiveDeps := b.recursive.dependencies(subq)
 	subqDirectDeps := b.direct.dependencies(subq)
 
-	var tablesToKeep []TableSet
+	var tablesToKeep MutableTableSet
 	sco := currScope
 	for sco != nil {
 		for _, table := range sco.tables {
-			tablesToKeep = append(tablesToKeep, table.getTableSet(b.org))
+			tablesToKeep.MergeInPlace(table.getTableSet(b.org))
 		}
 		sco = sco.parent
 	}
 
-	keep := MergeTableSets(tablesToKeep...)
+	keep := tablesToKeep.ToImmutable()
 	b.recursive[subq] = subqRecursiveDeps.KeepOnly(keep)
 	b.direct[subq] = subqDirectDeps.KeepOnly(keep)
 	return nil

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

@@ -245,45 +245,3 @@ func Single(bit int) Bitset {
 		return toBitset(words)
 	}
 }
-
-// Merge takes multiple Bitsets and returns a single Bitset that is the
-// bitwise-OR of all of them. It attempts to reduce allocations by creating
-// a single backing slice and merging all Bitsets into it.
-func Merge(bitsets ...Bitset) Bitset {
-	if len(bitsets) == 0 {
-		return ""
-	}
-	if len(bitsets) == 1 {
-		return bitsets[0]
-	}
-
-	// 1. Find the largest length among all input Bitsets.
-	maxLen := 0
-	for _, bs := range bitsets {
-		if len(bs) > maxLen {
-			maxLen = len(bs)
-		}
-	}
-
-	// 2. Allocate a single byte slice of size maxLen.
-	merged := make([]byte, maxLen)
-
-	// 3. Merge each Bitset into the merged buffer in-place.
-	for _, bs := range bitsets {
-		for i := 0; i < len(bs); i++ {
-			merged[i] |= bs[i]
-		}
-	}
-
-	// 4. Trim trailing zeros to avoid unnecessary storage.
-	trimmed := maxLen
-	for trimmed > 0 && merged[trimmed-1] == 0 {
-		trimmed--
-	}
-	if trimmed == 0 {
-		return "" // all bits are zero
-	}
-
-	merged = merged[:trimmed]
-	return toBitset(merged)
-}

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

@@ -22,11 +22,6 @@ type Mutable struct {
 	data []byte
 }
 
-// NewMutable creates an initially empty mutable BitSet.
-func NewMutable() *Mutable {
-	return &Mutable{}
-}
-
 // 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.
@@ -47,10 +42,5 @@ func (m *Mutable) AsImmutable() Bitset {
 	for trim > 0 && m.data[trim-1] == 0 {
 		trim--
 	}
-	if trim == 0 {
-		// Means all bits are zero -> empty
-		return ""
-	}
-
 	return toBitset(m.data[:trim])
 }

go/vt/vtgate/semantics/cte_table.go

@@ -167,16 +167,16 @@ func (cte *CTE) recursive(org originable) TableSet {
 	if cte.recursiveDeps != nil {
 		return *cte.recursiveDeps
 	}
-	var tables []TableSet
+	var tables 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
 	_ = sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
 		ate, ok := node.(*sqlparser.AliasedTableExpr)
 		if !ok {
 			return true, nil
 		}
-		tables = append(tables, org.tableSetFor(ate))
+		tables.MergeInPlace(org.tableSetFor(ate))
 		return true, nil
 	}, cte.Query)
-	return MergeTableSets(tables...)
+	return tables.ToImmutable()
 }

go/vt/vtgate/semantics/derived_table.go

@@ -92,15 +92,15 @@ func handleAliasedExpr(vTbl *DerivedTable, expr *sqlparser.AliasedExpr, cols sql
 }
 
 func handleUnexpandedStarExpression(tables []TableInfo, vTbl *DerivedTable, org originable) {
-	var tableSets []TableSet
+	var tableSets MutableTableSet
 	for _, table := range tables {
 		ts := table.getTableSet(org)
-		tableSets = append(tableSets, ts)
+		tableSets.MergeInPlace(ts)
 		if !table.authoritative() {
 			vTbl.isAuthoritative = false
 		}
 	}
-	vTbl.tables = MergeTableSets(tableSets...)
+	vTbl.tables = tableSets.ToImmutable()
 }
 
 // dependencies implements the TableInfo interface

go/vt/vtgate/semantics/scoper.go

@@ -228,11 +228,11 @@ func (s *scoper) up(cursor *sqlparser.Cursor) error {
 			s.popScope()
 		}
 	case *sqlparser.Select, *sqlparser.GroupBy, *sqlparser.Update, *sqlparser.Insert, *sqlparser.Union, *sqlparser.Delete:
-		var tables []TableSet
+		var tables MutableTableSet
 		for _, tableInfo := range s.currentScope().tables {
-			tables = append(tables, tableInfo.getTableSet(s.org))
+			tables.MergeInPlace(tableInfo.getTableSet(s.org))
 		}
-		s.statementIDs[s.currentScope().stmt] = MergeTableSets(tables...)
+		s.statementIDs[s.currentScope().stmt] = tables.ToImmutable()
 		s.popScope()
 	case *sqlparser.Where:
 		if node.Type != sqlparser.HavingClause {

go/vt/vtgate/semantics/semantic_table.go

@@ -19,8 +19,6 @@ package semantics
 import (
 	"fmt"
 
-	"vitess.io/vitess/go/vt/vtgate/semantics/bitset"
-
 	"vitess.io/vitess/go/mysql/collations"
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/key"
@@ -708,7 +706,7 @@ func (d ExprDependencies) dependencies(expr sqlparser.Expr) (deps TableSet) {
 		}()
 	}
 
-	depsCalc := bitset.NewMutable()
+	var depsCalc MutableTableSet
 	// During the original semantic analysis, all ColNames were found and bound to the corresponding tables
 	// Here, we'll walk the expression tree and look to see if we can find any sub-expressions
 	// that have already set dependencies.
@@ -722,15 +720,14 @@ func (d ExprDependencies) dependencies(expr sqlparser.Expr) (deps TableSet) {
 
 		set, found := d[expr]
 		if found {
-			depsCalc.Or(bitset.Bitset(set))
+			depsCalc.MergeInPlace(set)
 		}
 
 		// if we found a cached value, there is no need to continue down to visit children
 		return !found, nil
 	}, expr)
 
-	deps = TableSet(depsCalc.AsImmutable())
-
+	deps = depsCalc.ToImmutable()
 	return
 }
 

go/vt/vtgate/semantics/table_collector.go

@@ -169,7 +169,7 @@ func (tc *tableCollector) visitUnion(union *sqlparser.Union) error {
 	}
 
 	size := firstSelect.GetColumnCount()
-	recursiveDeps := make([][]TableSet, size)
+	recursiveDeps := make([]MutableTableSet, size)
 	typers := make([]evalengine.TypeAggregator, size)
 	collations := tc.org.collationEnv()
 
@@ -180,7 +180,7 @@ func (tc *tableCollector) visitUnion(union *sqlparser.Union) error {
 				continue
 			}
 			_, deps, qt := tc.org.depsForExpr(ae.Expr)
-			recursiveDeps[i] = append(recursiveDeps[i], deps)
+			recursiveDeps[i].MergeInPlace(deps)
 			if err := typers[i].Add(qt, collations); err != nil {
 				return err
 			}
@@ -194,7 +194,7 @@ func (tc *tableCollector) visitUnion(union *sqlparser.Union) error {
 	info.recursive = make([]TableSet, size)
 	for i, ts := range typers {
 		info.types = append(info.types, ts.Type())
-		info.recursive[i] = MergeTableSets(recursiveDeps[i]...)
+		info.recursive[i] = recursiveDeps[i].ToImmutable()
 	}
 	tc.unionInfo[union] = info
 	return nil