Skip to content

Commit

Permalink
feat: add planner support for simpler concatenate code when possible
Browse files Browse the repository at this point in the history 
Signed-off-by: Andres Taylor <[email protected]>
  • Loading branch information
@systay
systay committed Mar 19, 2024
1 parent 88e0edb commit 661fe19
Show file tree
Hide file tree
Showing 8 changed files with 142 additions and 32 deletions.
8 changes: 6 additions & 2 deletions go/vt/vtgate/engine/coerce.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ package engine

import (
"context"
"fmt"

"vitess.io/vitess/go/sqltypes"
querypb "vitess.io/vitess/go/vt/proto/query"
Expand Down Expand Up @@ -151,8 +152,11 @@ func (c *Coerce) Inputs() ([]Primitive, []map[string]any) {

func (c *Coerce) description() PrimitiveDescription {
var cols []string
for _, typ := range c.Types {
cols = append(cols, typ.Type().String())
for idx, typ := range c.Types {
if typ == nil {
continue
}
cols = append(cols, fmt.Sprintf("%d:%s", idx, typ.Type().String()))
}
return PrimitiveDescription{
OperatorType: "Coerce",
Expand Down
8 changes: 8 additions & 0 deletions go/vt/vtgate/planbuilder/concatenate.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import (

type concatenate struct {
sources []logicalPlan
coerced bool

// These column offsets do not need to be typed checked - they usually contain weight_string()
// columns that are not going to be returned to the user
Expand All @@ -37,5 +38,12 @@ func (c *concatenate) Primitive() engine.Primitive {
sources = append(sources, source.Primitive())
}

if c.coerced {
// types are already handled, let's use the fast concatenate
return &engine.SimpleConcatenate{
Sources: sources,
}
}

return engine.NewConcatenate(sources, c.noNeedToTypeCheck)
}
14 changes: 14 additions & 0 deletions go/vt/vtgate/planbuilder/filter.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ package planbuilder

import (
"vitess.io/vitess/go/vt/vtgate/engine"
"vitess.io/vitess/go/vt/vtgate/evalengine"
)

type (
Expand All @@ -35,3 +36,16 @@ func (l *filter) Primitive() engine.Primitive {
l.efilter.Input = l.input.Primitive()
return l.efilter
}

type coercePlan struct {
input logicalPlan
columns []*evalengine.Type
}

func (c *coercePlan) Primitive() engine.Primitive {
src := c.input.Primitive()
return &engine.Coerce{
Source: src,
Types: c.columns,
}
}
100 changes: 92 additions & 8 deletions go/vt/vtgate/planbuilder/operator_transformers.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -811,25 +811,109 @@ func getAllTableNames(op *operators.Route) ([]string, error) {
}

func transformUnionPlan(ctx *plancontext.PlanningContext, op *operators.Union) (logicalPlan, error) {
sources, err := slice.MapWithError(op.Sources, func(src operators.Operator) (logicalPlan, error) {
sources, coerced, err := coercedInputs(ctx, op)
if err != nil {
return nil, err
}

return &concatenate{
sources: sources,
coerced: coerced,
}, nil
}

func typeForExpr(ctx *plancontext.PlanningContext, e sqlparser.Expr) (evalengine.Type, bool) {
if typ, found := ctx.SemTable.TypeForExpr(e); found {
return typ, true
}

cfg := &evalengine.Config{
ResolveColumn: func(name *sqlparser.ColName) (int, error) {
return 0, nil // we are not going to use these for anything other than getting the type
},
ResolveType: func(expr sqlparser.Expr) (evalengine.Type, bool) {
return ctx.SemTable.TypeForExpr(e)
},
Collation: ctx.SemTable.Collation,
Environment: ctx.VSchema.Environment(),
}
evalExpr, err := evalengine.Translate(e, cfg)
if err != nil {
return evalengine.Type{}, false
}
env := evalengine.ExpressionEnv{
BindVars: nil,
Row: nil,
Fields: nil,
}
typ, err := env.TypeOf(evalExpr)
if err != nil {
return evalengine.Type{}, false
}
ctx.SemTable.ExprTypes[e] = typ
return typ, true
}

func coercedInputs(ctx *plancontext.PlanningContext, op *operators.Union) ([]logicalPlan, bool, error) {
orgSources, err := slice.MapWithError(op.Sources, func(src operators.Operator) (logicalPlan, error) {
plan, err := transformToLogicalPlan(ctx, src)
if err != nil {
return nil, err
}
return plan, nil
})
if err != nil {
return nil, err
return nil, false, err
}
collationEnv := ctx.VSchema.Environment().CollationEnv()
typers := make([]evalengine.TypeAggregator, len(op.Sources[0].GetColumns(ctx)))
for _, src := range op.Sources {
cols := src.GetColumns(ctx)
for idx, col := range cols {
typ, found := typeForExpr(ctx, col.Expr)
if !found {
return orgSources, false, nil
}
err := typers[idx].Add(typ, collationEnv)
if err != nil {
// let's ignore this and just return the
return orgSources, false, nil
}
}
}

if len(sources) == 1 {
return sources[0], nil
newSources := make([]logicalPlan, 0, len(orgSources))
for srcIdx, src := range op.Sources {
cols := src.GetColumns(ctx)
coerceTypes := make([]*evalengine.Type, len(cols))
coerce := false
for colIdx, col := range cols {
typ, found := ctx.SemTable.TypeForExpr(col.Expr)
if !found {
return orgSources, false, nil
}
resultType := typers[colIdx].Type()
if resultType.Type() == sqltypes.Unknown {
// if the resulting type is a null type, the type aggregator probably messed up the
// type calculus, and we can't trust these types
return orgSources, false, nil
}
if resultType != typ {
coerceTypes[colIdx] = &resultType
coerce = true
}
}
if coerce {
newSources = append(newSources, &coercePlan{
input: orgSources[srcIdx],
columns: coerceTypes,
})
} else {
newSources = append(newSources, orgSources[srcIdx])
}
}
return &concatenate{
sources: sources,
noNeedToTypeCheck: nil,
}, nil

return newSources, true, nil
}

func transformLimit(ctx *plancontext.PlanningContext, op *operators.Limit) (logicalPlan, error) {
Expand Down
2 changes: 1 addition & 1 deletion go/vt/vtgate/planbuilder/testdata/cte_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1901,7 +1901,7 @@
],
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Projection",
Expand Down
18 changes: 9 additions & 9 deletions go/vt/vtgate/planbuilder/testdata/info_schema57_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@
],
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -170,7 +170,7 @@
],
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -788,7 +788,7 @@
"Aggregates": "sum(0) AS sum(found)",
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -827,7 +827,7 @@
"QueryType": "SELECT",
"Original": "select found from (select 1 as found from information_schema.`tables` where table_schema = 'music' union all (select 1 as found from information_schema.views where table_schema = 'music' limit 1)) as t",
"Instructions": {
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -864,7 +864,7 @@
"QueryType": "SELECT",
"Original": "select 1 as found from information_schema.`tables` where table_schema = 'music' and table_schema = 'Music' union all (select 1 as found from information_schema.views where table_schema = 'music' and table_schema = 'user' limit 1)",
"Instructions": {
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -901,7 +901,7 @@
"QueryType": "SELECT",
"Original": "select 1 as found from information_schema.`tables` where table_schema = 'music' and table_schema = 'Music' union all (select 1 as found from information_schema.views where table_schema = 'music' and table_schema = 'user' limit 1)",
"Instructions": {
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -946,7 +946,7 @@
"Inputs": [
{
"InputName": "SubQuery",
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -1004,7 +1004,7 @@
],
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -1057,7 +1057,7 @@
],
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down
16 changes: 8 additions & 8 deletions go/vt/vtgate/planbuilder/testdata/info_schema80_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@
],
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -854,7 +854,7 @@
"Aggregates": "sum(0) AS sum(found)",
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -893,7 +893,7 @@
"QueryType": "SELECT",
"Original": "select found from (select 1 as found from information_schema.`tables` where table_schema = 'music' union all (select 1 as found from information_schema.views where table_schema = 'music' limit 1)) as t",
"Instructions": {
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -930,7 +930,7 @@
"QueryType": "SELECT",
"Original": "select 1 as found from information_schema.`tables` where table_schema = 'music' and table_schema = 'Music' union all (select 1 as found from information_schema.views where table_schema = 'music' and table_schema = 'user' limit 1)",
"Instructions": {
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -967,7 +967,7 @@
"QueryType": "SELECT",
"Original": "select 1 as found from information_schema.`tables` where table_schema = 'music' and table_schema = 'Music' union all (select 1 as found from information_schema.views where table_schema = 'music' and table_schema = 'user' limit 1)",
"Instructions": {
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -1012,7 +1012,7 @@
"Inputs": [
{
"InputName": "SubQuery",
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -1070,7 +1070,7 @@
],
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -1123,7 +1123,7 @@
],
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down
8 changes: 4 additions & 4 deletions go/vt/vtgate/planbuilder/testdata/union_cases.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -443,7 +443,7 @@
],
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -696,7 +696,7 @@
],
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Aggregate",
Expand Down Expand Up @@ -1144,7 +1144,7 @@
],
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down Expand Up @@ -1235,7 +1235,7 @@
],
"Inputs": [
{
"OperatorType": "Concatenate",
"OperatorType": "SimpleConcatenate",
"Inputs": [
{
"OperatorType": "Route",
Expand Down

0 comments on commit 661fe19

Please sign in to comment.