From cf6908da84ae0babc43c29131dfafc3671074d75 Mon Sep 17 00:00:00 2001 From: Vicent Marti Date: Fri, 17 Jan 2025 17:59:13 +0100 Subject: [PATCH] wip: fewer allocations Signed-off-by: Vicent Marti --- .../vtgate/planbuilder/operators/ast_to_op.go | 7 ++-- go/vt/vtgate/planbuilder/operators/helpers.go | 6 +-- .../planbuilder/operators/querygraph.go | 6 +-- go/vt/vtgate/planbuilder/operators/route.go | 6 +-- .../operators/subquery_planning.go | 10 ++--- go/vt/vtgate/semantics/analyzer.go | 2 +- go/vt/vtgate/semantics/binder.go | 24 +++++------ go/vt/vtgate/semantics/bitset/bitset.go | 42 ------------------- go/vt/vtgate/semantics/bitset/mutable.go | 10 ----- go/vt/vtgate/semantics/cte_table.go | 6 +-- go/vt/vtgate/semantics/derived_table.go | 6 +-- go/vt/vtgate/semantics/scoper.go | 6 +-- go/vt/vtgate/semantics/semantic_table.go | 9 ++-- go/vt/vtgate/semantics/table_collector.go | 6 +-- go/vt/vtgate/semantics/table_set.go | 21 ++++------ go/vt/vtgate/semantics/vtable.go | 6 +-- 16 files changed, 56 insertions(+), 117 deletions(-) diff --git a/go/vt/vtgate/planbuilder/operators/ast_to_op.go b/go/vt/vtgate/planbuilder/operators/ast_to_op.go index 8f9246f60b8..da6a7688d46 100644 --- a/go/vt/vtgate/planbuilder/operators/ast_to_op.go +++ b/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 diff --git a/go/vt/vtgate/planbuilder/operators/helpers.go b/go/vt/vtgate/planbuilder/operators/helpers.go index 10b616c7217..af8e2607b18 100644 --- a/go/vt/vtgate/planbuilder/operators/helpers.go +++ b/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 diff --git a/go/vt/vtgate/planbuilder/operators/querygraph.go b/go/vt/vtgate/planbuilder/operators/querygraph.go index 9abe74a0fec..ac885383801 100644 --- a/go/vt/vtgate/planbuilder/operators/querygraph.go +++ b/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 diff --git a/go/vt/vtgate/planbuilder/operators/route.go b/go/vt/vtgate/planbuilder/operators/route.go index a22782e2945..a459c5a3f7a 100644 --- a/go/vt/vtgate/planbuilder/operators/route.go +++ b/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() } diff --git a/go/vt/vtgate/planbuilder/operators/subquery_planning.go b/go/vt/vtgate/planbuilder/operators/subquery_planning.go index 01dad5f1327..764b6439ad6 100644 --- a/go/vt/vtgate/planbuilder/operators/subquery_planning.go +++ b/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 diff --git a/go/vt/vtgate/semantics/analyzer.go b/go/vt/vtgate/semantics/analyzer.go index c4e7dc55866..32b3e8930ac 100644 --- a/go/vt/vtgate/semantics/analyzer.go +++ b/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, diff --git a/go/vt/vtgate/semantics/binder.go b/go/vt/vtgate/semantics/binder.go index 6883828a8a2..f0a76acabe3 100644 --- a/go/vt/vtgate/semantics/binder.go +++ b/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 diff --git a/go/vt/vtgate/semantics/bitset/bitset.go b/go/vt/vtgate/semantics/bitset/bitset.go index b834b158e67..898d55e1d95 100644 --- a/go/vt/vtgate/semantics/bitset/bitset.go +++ b/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) -} diff --git a/go/vt/vtgate/semantics/bitset/mutable.go b/go/vt/vtgate/semantics/bitset/mutable.go index 2bf42672e07..7782ad4224b 100644 --- a/go/vt/vtgate/semantics/bitset/mutable.go +++ b/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]) } diff --git a/go/vt/vtgate/semantics/cte_table.go b/go/vt/vtgate/semantics/cte_table.go index 7703d320046..c164069ab38 100644 --- a/go/vt/vtgate/semantics/cte_table.go +++ b/go/vt/vtgate/semantics/cte_table.go @@ -167,7 +167,7 @@ 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) { @@ -175,8 +175,8 @@ func (cte *CTE) recursive(org originable) TableSet { 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() } diff --git a/go/vt/vtgate/semantics/derived_table.go b/go/vt/vtgate/semantics/derived_table.go index c4b71fc1a35..38169514cd9 100644 --- a/go/vt/vtgate/semantics/derived_table.go +++ b/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 diff --git a/go/vt/vtgate/semantics/scoper.go b/go/vt/vtgate/semantics/scoper.go index 2cc5e4c3e97..ee91b89d42f 100644 --- a/go/vt/vtgate/semantics/scoper.go +++ b/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 { diff --git a/go/vt/vtgate/semantics/semantic_table.go b/go/vt/vtgate/semantics/semantic_table.go index 96b8a2e8ea1..febadb797d5 100644 --- a/go/vt/vtgate/semantics/semantic_table.go +++ b/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 } diff --git a/go/vt/vtgate/semantics/table_collector.go b/go/vt/vtgate/semantics/table_collector.go index f6c8bb4d21c..5edf350dc08 100644 --- a/go/vt/vtgate/semantics/table_collector.go +++ b/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 diff --git a/go/vt/vtgate/semantics/table_set.go b/go/vt/vtgate/semantics/table_set.go index eb2a386d55b..ff3ca2fb7e6 100644 --- a/go/vt/vtgate/semantics/table_set.go +++ b/go/vt/vtgate/semantics/table_set.go @@ -18,7 +18,6 @@ package semantics import ( "fmt" - "unsafe" "vitess.io/vitess/go/vt/vtgate/semantics/bitset" ) @@ -115,18 +114,14 @@ func EmptyTableSet() TableSet { return "" } -// MergeTableSets merges all the given TableSet into a single one -func MergeTableSets(tableSets ...TableSet) TableSet { - if len(tableSets) == 0 { - return "" - } +type MutableTableSet struct { + bitset bitset.Mutable +} - // Trick: re-interpret slice header of tableSets as []Bitset - // This is safe because the memory layout of []TableSet and []Bitset is the same - // The alternative would be to loop over all TableSets and convert them to Bitsets, but that would be slower - bs := *(*[]bitset.Bitset)(unsafe.Pointer(&tableSets)) +func (ts *MutableTableSet) MergeInPlace(other TableSet) { + ts.bitset.Or(bitset.Bitset(other)) +} - // Now pass it to MergeMany - merged := bitset.Merge(bs...) - return TableSet(merged) +func (ts *MutableTableSet) ToImmutable() TableSet { + return TableSet(ts.bitset.AsImmutable()) } diff --git a/go/vt/vtgate/semantics/vtable.go b/go/vt/vtgate/semantics/vtable.go index 512c1601f11..a6c8a3eff8c 100644 --- a/go/vt/vtgate/semantics/vtable.go +++ b/go/vt/vtgate/semantics/vtable.go @@ -165,14 +165,14 @@ func selectExprsToInfos( colNames = append(colNames, expr.As.String()) } case *sqlparser.StarExpr: - var tableSets []TableSet + var tableSets MutableTableSet for _, table := range tables { - tableSets = append(tableSets, table.getTableSet(org)) + tableSets.MergeInPlace(table.getTableSet(org)) if !table.authoritative() { isAuthoritative = false } } - ts = MergeTableSets(tableSets...) + ts = tableSets.ToImmutable() } } return