From 661fe19b97d5d5e47d724d9c4b8512d219eda3cf Mon Sep 17 00:00:00 2001 From: Andres Taylor Date: Tue, 19 Mar 2024 07:52:54 +0100 Subject: [PATCH] feat: add planner support for simpler concatenate code when possible Signed-off-by: Andres Taylor --- go/vt/vtgate/engine/coerce.go | 8 +- go/vt/vtgate/planbuilder/concatenate.go | 8 ++ go/vt/vtgate/planbuilder/filter.go | 14 +++ .../planbuilder/operator_transformers.go | 100 ++++++++++++++++-- .../planbuilder/testdata/cte_cases.json | 2 +- .../testdata/info_schema57_cases.json | 18 ++-- .../testdata/info_schema80_cases.json | 16 +-- .../planbuilder/testdata/union_cases.json | 8 +- 8 files changed, 142 insertions(+), 32 deletions(-) diff --git a/go/vt/vtgate/engine/coerce.go b/go/vt/vtgate/engine/coerce.go index e4ee2184a33..cd707b304ae 100644 --- a/go/vt/vtgate/engine/coerce.go +++ b/go/vt/vtgate/engine/coerce.go @@ -18,6 +18,7 @@ package engine import ( "context" + "fmt" "vitess.io/vitess/go/sqltypes" querypb "vitess.io/vitess/go/vt/proto/query" @@ -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", diff --git a/go/vt/vtgate/planbuilder/concatenate.go b/go/vt/vtgate/planbuilder/concatenate.go index 81cbe3d5b65..5c512ee4c1b 100644 --- a/go/vt/vtgate/planbuilder/concatenate.go +++ b/go/vt/vtgate/planbuilder/concatenate.go @@ -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 @@ -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) } diff --git a/go/vt/vtgate/planbuilder/filter.go b/go/vt/vtgate/planbuilder/filter.go index c3686380446..45b7f2affba 100644 --- a/go/vt/vtgate/planbuilder/filter.go +++ b/go/vt/vtgate/planbuilder/filter.go @@ -18,6 +18,7 @@ package planbuilder import ( "vitess.io/vitess/go/vt/vtgate/engine" + "vitess.io/vitess/go/vt/vtgate/evalengine" ) type ( @@ -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, + } +} diff --git a/go/vt/vtgate/planbuilder/operator_transformers.go b/go/vt/vtgate/planbuilder/operator_transformers.go index 6d2c2317517..3ade6931b27 100644 --- a/go/vt/vtgate/planbuilder/operator_transformers.go +++ b/go/vt/vtgate/planbuilder/operator_transformers.go @@ -811,7 +811,51 @@ 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 @@ -819,17 +863,57 @@ func transformUnionPlan(ctx *plancontext.PlanningContext, op *operators.Union) ( 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) { diff --git a/go/vt/vtgate/planbuilder/testdata/cte_cases.json b/go/vt/vtgate/planbuilder/testdata/cte_cases.json index 0d7d9020ac2..80d3ef6e57a 100644 --- a/go/vt/vtgate/planbuilder/testdata/cte_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/cte_cases.json @@ -1901,7 +1901,7 @@ ], "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Projection", diff --git a/go/vt/vtgate/planbuilder/testdata/info_schema57_cases.json b/go/vt/vtgate/planbuilder/testdata/info_schema57_cases.json index 1727e372490..0a11d83dc02 100644 --- a/go/vt/vtgate/planbuilder/testdata/info_schema57_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/info_schema57_cases.json @@ -107,7 +107,7 @@ ], "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", @@ -170,7 +170,7 @@ ], "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", @@ -788,7 +788,7 @@ "Aggregates": "sum(0) AS sum(found)", "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", @@ -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", @@ -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", @@ -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", @@ -946,7 +946,7 @@ "Inputs": [ { "InputName": "SubQuery", - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", @@ -1004,7 +1004,7 @@ ], "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", @@ -1057,7 +1057,7 @@ ], "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", diff --git a/go/vt/vtgate/planbuilder/testdata/info_schema80_cases.json b/go/vt/vtgate/planbuilder/testdata/info_schema80_cases.json index d0a4911fb74..12f143d7de5 100644 --- a/go/vt/vtgate/planbuilder/testdata/info_schema80_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/info_schema80_cases.json @@ -107,7 +107,7 @@ ], "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", @@ -854,7 +854,7 @@ "Aggregates": "sum(0) AS sum(found)", "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", @@ -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", @@ -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", @@ -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", @@ -1012,7 +1012,7 @@ "Inputs": [ { "InputName": "SubQuery", - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", @@ -1070,7 +1070,7 @@ ], "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", @@ -1123,7 +1123,7 @@ ], "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", diff --git a/go/vt/vtgate/planbuilder/testdata/union_cases.json b/go/vt/vtgate/planbuilder/testdata/union_cases.json index 76f1fa460ca..d5920a3cbbe 100644 --- a/go/vt/vtgate/planbuilder/testdata/union_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/union_cases.json @@ -443,7 +443,7 @@ ], "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", @@ -696,7 +696,7 @@ ], "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Aggregate", @@ -1144,7 +1144,7 @@ ], "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route", @@ -1235,7 +1235,7 @@ ], "Inputs": [ { - "OperatorType": "Concatenate", + "OperatorType": "SimpleConcatenate", "Inputs": [ { "OperatorType": "Route",