semantics: Use a BitSet (#11819) (#11837)

* make sure planner works well with more than 64 tables

Signed-off-by: Andres Taylor <andres@planetscale.com>

* semantics: use an immutable bitset

Signed-off-by: Vicent Marti <vmg@strn.cat>

* bitset: add documentation

Signed-off-by: Vicent Marti <vmg@strn.cat>

* bitset: add license

Signed-off-by: Vicent Marti <vmg@strn.cat>

* use IsEmpty and NonEmpty more

Signed-off-by: Andres Taylor <andres@planetscale.com>

* rename file

Signed-off-by: Andres Taylor <andres@planetscale.com>

* bitset: better documentation

Signed-off-by: Vicent Marti <vmg@strn.cat>

Signed-off-by: Andres Taylor <andres@planetscale.com>
Signed-off-by: Vicent Marti <vmg@strn.cat>
Co-authored-by: Andres Taylor <andres@planetscale.com>
Signed-off-by: Manan Gupta <manan@planetscale.com>

Signed-off-by: Andres Taylor <andres@planetscale.com>
Signed-off-by: Vicent Marti <vmg@strn.cat>
Signed-off-by: Manan Gupta <manan@planetscale.com>
Co-authored-by: Vicent Martí <42793+vmg@users.noreply.github.com>
Co-authored-by: Andres Taylor <andres@planetscale.com>

go/vt/vtgate/planbuilder/abstract/concatenate.go

@@ -40,7 +40,7 @@ func (*Concatenate) iLogical() {}
 func (c *Concatenate) TableID() semantics.TableSet {
  var tableSet semantics.TableSet
  for _, source := range c.Sources {
- tableSet.MergeInPlace(source.TableID())
+ tableSet = tableSet.Merge(source.TableID())
  }
  return tableSet
 }

go/vt/vtgate/planbuilder/concatenateGen4.go

@@ -104,7 +104,7 @@ func (c *concatenateGen4) Rewrite(inputs ...logicalPlan) error {
 func (c *concatenateGen4) ContainsTables() semantics.TableSet {
  var tableSet semantics.TableSet
  for _, source := range c.sources {
- tableSet.MergeInPlace(source.ContainsTables())
+ tableSet = tableSet.Merge(source.ContainsTables())
  }
  return tableSet
 }

go/vt/vtgate/planbuilder/horizon_planning.go

@@ -204,7 +204,7 @@ func checkIfAlreadyExists(expr *sqlparser.AliasedExpr, node sqlparser.SelectStat
  selectExprDep := semTable.RecursiveDeps(selectExpr.Expr)
 
  // Check that the two expressions have the same dependencies
- if !selectExprDep.Equals(exprDep) {
+ if selectExprDep != exprDep {
  continue
  }
 

go/vt/vtgate/planbuilder/physical/operator_funcs.go

@@ -345,7 +345,7 @@ func BreakExpressionInLHSandRHS(
  switch node := cursor.Node().(type) {
  case *sqlparser.ColName:
  deps := ctx.SemTable.RecursiveDeps(node)
- if deps.NumberOfTables() == 0 {
+ if deps.IsEmpty() {
  err = vterrors.Errorf(vtrpcpb.Code_INTERNAL, "unknown column. has the AST been copied?")
  return false
  }

go/vt/vtgate/planbuilder/physical/subquery_planning.go

@@ -84,21 +84,6 @@ func optimizeSubQuery(ctx *plancontext.PlanningContext, op *abstract.SubQuery) (
 
 func mergeSubQueryOp(ctx *plancontext.PlanningContext, outer *Route, inner *Route, subq *abstract.SubQueryInner) (*Route, error) {
  subq.ExtractedSubquery.NeedsRewrite = true
-
- // go over the subquery and add its tables to the one's solved by the route it is merged with
- // this is needed to so that later when we try to push projections, we get the correct
- // solved tableID from the route, since it also includes the tables from the subquery after merging
- err := sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
- switch n := node.(type) {
- case *sqlparser.AliasedTableExpr:
- ts := outer.TableID()
- ts.MergeInPlace(ctx.SemTable.TableSetFor(n))
- }
- return true, nil
- }, subq.ExtractedSubquery.Subquery)
- if err != nil {
- return nil, err
- }
  outer.SysTableTableSchema = append(outer.SysTableTableSchema, inner.SysTableTableSchema...)
  for k, v := range inner.SysTableTableName {
  if outer.SysTableTableName == nil {
@@ -126,7 +111,7 @@ func mergeSubQueryOp(ctx *plancontext.PlanningContext, outer *Route, inner *Rout
  }
  }
 
- err = outer.resetRoutingSelections(ctx)
+ err := outer.resetRoutingSelections(ctx)
  if err != nil {
  return nil, err
  }
@@ -353,11 +338,9 @@ func rewriteColumnsInSubqueryOpForJoin(
 
  // update the dependencies for the subquery by removing the dependencies from the innerOp
  tableSet := ctx.SemTable.Direct[subQueryInner.ExtractedSubquery.Subquery]
- tableSet.RemoveInPlace(resultInnerOp.TableID())
- ctx.SemTable.Direct[subQueryInner.ExtractedSubquery.Subquery] = tableSet
+ ctx.SemTable.Direct[subQueryInner.ExtractedSubquery.Subquery] = tableSet.Remove(resultInnerOp.TableID())
  tableSet = ctx.SemTable.Recursive[subQueryInner.ExtractedSubquery.Subquery]
- tableSet.RemoveInPlace(resultInnerOp.TableID())
- ctx.SemTable.Recursive[subQueryInner.ExtractedSubquery.Subquery] = tableSet
+ ctx.SemTable.Recursive[subQueryInner.ExtractedSubquery.Subquery] = tableSet.Remove(resultInnerOp.TableID())
 
  // return any error while rewriting
  return resultInnerOp, rewriteError
@@ -388,7 +371,7 @@ func createCorrelatedSubqueryOp(
  // we do so by checking that the column names are the same and their recursive dependencies are the same
  // so if the column names user.a and a would also be equal if the latter is also referencing the user table
  for colName, bindVar := range bindVars {
- if node.Name.Equal(colName.Name) && ctx.SemTable.RecursiveDeps(node).Equals(ctx.SemTable.RecursiveDeps(colName)) {
+ if node.Name.Equal(colName.Name) && (ctx.SemTable.RecursiveDeps(node) == ctx.SemTable.RecursiveDeps(colName)) {
  cursor.Replace(sqlparser.NewArgument(bindVar))
  return false
  }

go/vt/vtgate/planbuilder/physical/union.go

@@ -37,7 +37,7 @@ var _ abstract.PhysicalOperator = (*Union)(nil)
 func (u *Union) TableID() semantics.TableSet {
  ts := semantics.EmptyTableSet()
  for _, source := range u.Sources {
- ts.MergeInPlace(source.TableID())
+ ts = ts.Merge(source.TableID())
  }
  return ts
 }

go/vt/vtgate/planbuilder/plan_test.go

@@ -243,6 +243,7 @@ func TestPlan(t *testing.T) {
  testFile(t, "use_cases.json", testOutputTempDir, vschemaWrapper, false)
  testFile(t, "set_cases.json", testOutputTempDir, vschemaWrapper, false)
  testFile(t, "union_cases.json", testOutputTempDir, vschemaWrapper, false)
+ testFile(t, "large_union_cases.json", testOutputTempDir, vschemaWrapper, false)
  testFile(t, "transaction_cases.json", testOutputTempDir, vschemaWrapper, false)
  testFile(t, "lock_cases.json", testOutputTempDir, vschemaWrapper, false)
  testFile(t, "large_cases.json", testOutputTempDir, vschemaWrapper, false)

go/vt/vtgate/planbuilder/projection_pushing.go

@@ -76,7 +76,7 @@ func pushProjectionIntoConcatenate(ctx *plancontext.PlanningContext, expr *sqlpa
  if err != nil {
  return 0, false, err
  }
- if added && ctx.SemTable.DirectDeps(expr.Expr).NumberOfTables() > 0 {
+ if added && ctx.SemTable.DirectDeps(expr.Expr).NonEmpty() {
  return 0, false, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "pushing projection %v on concatenate should reference an existing column", sqlparser.String(expr))
  }
  if added {