diff --git a/go/vt/vtgate/planbuilder/operators/union.go b/go/vt/vtgate/planbuilder/operators/union.go index 8fd06adb219..0731d8c888e 100644 --- a/go/vt/vtgate/planbuilder/operators/union.go +++ b/go/vt/vtgate/planbuilder/operators/union.go @@ -19,6 +19,7 @@ package operators import ( "fmt" "slices" + "strings" "vitess.io/vitess/go/slice" "vitess.io/vitess/go/vt/sqlparser" @@ -101,19 +102,10 @@ func (u *Union) AddPredicate(ctx *plancontext.PlanningContext, expr sqlparser.Ex if !ok { panic(vterrors.VT12001("pushing predicates on UNION where the first SELECT contains * or NEXT")) } - offsets[ae.ColumnName()] = i - } - - if jp, ok := expr.(*predicates.JoinPredicate); ok { - expr = jp.Current() - ctx.PredTracker.Skip(jp.ID) - } - - needsFilter, exprPerSource := u.predicatePerSource(expr, offsets) - if needsFilter { - return newFilter(u, expr) + offsets[strings.ToLower(ae.ColumnName())] = i } + exprPerSource := u.predicatePerSource(ctx, expr, offsets) for i, src := range u.Sources { u.Sources[i] = src.AddPredicate(ctx, exprPerSource[i]) } @@ -121,12 +113,20 @@ func (u *Union) AddPredicate(ctx *plancontext.PlanningContext, expr sqlparser.Ex return u } -func (u *Union) predicatePerSource(expr sqlparser.Expr, offsets map[string]int) (bool, []sqlparser.Expr) { - needsFilter := false +func (u *Union) predicatePerSource(ctx *plancontext.PlanningContext, expr sqlparser.Expr, offsets map[string]int) []sqlparser.Expr { exprPerSource := make([]sqlparser.Expr, len(u.Sources)) for i := range u.Sources { - predicate := sqlparser.CopyOnRewrite(expr, nil, func(cursor *sqlparser.CopyOnWriteCursor) { + predicate := expr + + if jp, ok := predicate.(*predicates.JoinPredicate); ok { + // Create a new JoinPredicate for each source to keep tracking working + // We can't use `*JoinPredicate.Clone` here as that would update the tracker and overwrite + // the expression for the original predicate + predicate = ctx.PredTracker.NewJoinPredicate(jp.Current()) + } + + predicate = sqlparser.CopyOnRewrite(predicate, nil, func(cursor *sqlparser.CopyOnWriteCursor) { col, ok := cursor.Node().(*sqlparser.ColName) if !ok { return @@ -134,9 +134,7 @@ func (u *Union) predicatePerSource(expr sqlparser.Expr, offsets map[string]int) idx, ok := offsets[col.Name.Lowered()] if !ok { - needsFilter = true - cursor.StopTreeWalk() - return + panic(vterrors.VT13001(fmt.Sprintf("could not find the column '%s' on the UNION", sqlparser.String(col)))) } sel := u.GetSelectFor(i) @@ -151,7 +149,7 @@ func (u *Union) predicatePerSource(expr sqlparser.Expr, offsets map[string]int) exprPerSource[i] = predicate } - return needsFilter, exprPerSource + return exprPerSource } func (u *Union) GetSelectFor(source int) *sqlparser.Select { diff --git a/go/vt/vtgate/planbuilder/testdata/union_cases.json b/go/vt/vtgate/planbuilder/testdata/union_cases.json index 7555662b078..6fccfba3eb7 100644 --- a/go/vt/vtgate/planbuilder/testdata/union_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/union_cases.json @@ -1357,45 +1357,39 @@ "QueryType": "SELECT", "Original": "select * from (select kcu.`COLUMN_NAME` from `information_schema`.`key_column_usage` `kcu` where `kcu`.`table_schema` = 'user' and `kcu`.`table_name` = 'user_extra' union select kcu.`COLUMN_NAME` from `information_schema`.`key_column_usage` `kcu` where `kcu`.`table_schema` = 'user' and `kcu`.`table_name` = 'music') `kcu` where `COLUMN_NAME` = 'primary'", "Instructions": { - "OperatorType": "Filter", - "Predicate": "`COLUMN_NAME` = 'primary'", + "OperatorType": "Distinct", + "Collations": [ + "0: utf8mb3_general_ci" + ], "Inputs": [ { - "OperatorType": "Distinct", - "Collations": [ - "0: utf8mb3_general_ci" - ], + "OperatorType": "Concatenate", "Inputs": [ { - "OperatorType": "Concatenate", - "Inputs": [ - { - "OperatorType": "Route", - "Variant": "DBA", - "Keyspace": { - "Name": "main", - "Sharded": false - }, - "FieldQuery": "select kcu.`COLUMN_NAME` from information_schema.key_column_usage as kcu where 1 != 1", - "Query": "select distinct kcu.`COLUMN_NAME` from information_schema.key_column_usage as kcu where kcu.table_schema = :__vtschemaname /* VARCHAR */ and kcu.`table_name` = :kcu_table_name /* VARCHAR */", - "SysTableTableName": "[kcu_table_name:'user_extra']", - "SysTableTableSchema": "['user']", - "Table": "information_schema.key_column_usage" - }, - { - "OperatorType": "Route", - "Variant": "DBA", - "Keyspace": { - "Name": "main", - "Sharded": false - }, - "FieldQuery": "select kcu.`COLUMN_NAME` from information_schema.key_column_usage as kcu where 1 != 1", - "Query": "select distinct kcu.`COLUMN_NAME` from information_schema.key_column_usage as kcu where kcu.table_schema = :__vtschemaname /* VARCHAR */ and kcu.`table_name` = :kcu_table_name1 /* VARCHAR */", - "SysTableTableName": "[kcu_table_name1:'music']", - "SysTableTableSchema": "['user']", - "Table": "information_schema.key_column_usage" - } - ] + "OperatorType": "Route", + "Variant": "DBA", + "Keyspace": { + "Name": "main", + "Sharded": false + }, + "FieldQuery": "select kcu.`COLUMN_NAME` from information_schema.key_column_usage as kcu where 1 != 1", + "Query": "select distinct kcu.`COLUMN_NAME` from information_schema.key_column_usage as kcu where kcu.table_schema = :__vtschemaname /* VARCHAR */ and kcu.`table_name` = :kcu_table_name /* VARCHAR */ and `COLUMN_NAME` = 'primary'", + "SysTableTableName": "[kcu_table_name:'user_extra']", + "SysTableTableSchema": "['user']", + "Table": "information_schema.key_column_usage" + }, + { + "OperatorType": "Route", + "Variant": "DBA", + "Keyspace": { + "Name": "main", + "Sharded": false + }, + "FieldQuery": "select kcu.`COLUMN_NAME` from information_schema.key_column_usage as kcu where 1 != 1", + "Query": "select distinct kcu.`COLUMN_NAME` from information_schema.key_column_usage as kcu where kcu.table_schema = :__vtschemaname /* VARCHAR */ and kcu.`table_name` = :kcu_table_name1 /* VARCHAR */ and `COLUMN_NAME` = 'primary'", + "SysTableTableName": "[kcu_table_name1:'music']", + "SysTableTableSchema": "['user']", + "Table": "information_schema.key_column_usage" } ] } @@ -1439,56 +1433,50 @@ "QueryType": "SELECT", "Original": "select * from (select * from `information_schema`.`key_column_usage` `kcu` where `kcu`.`table_schema` = 'user' and `kcu`.`table_name` = 'user_extra' union select * from `information_schema`.`key_column_usage` `kcu` where `kcu`.`table_schema` = 'user' and `kcu`.`table_name` = 'music') `kcu` where `constraint_name` = 'primary'", "Instructions": { - "OperatorType": "Filter", - "Predicate": "`constraint_name` = 'primary'", + "OperatorType": "Distinct", + "Collations": [ + "0: utf8mb3_general_ci", + "1: utf8mb3_general_ci", + "2: utf8mb3_general_ci", + "3: utf8mb3_general_ci", + "4: utf8mb3_general_ci", + "5: utf8mb3_general_ci", + "6: utf8mb3_general_ci", + "7", + "8", + "9: utf8mb3_general_ci", + "10: utf8mb3_general_ci", + "11: utf8mb3_general_ci" + ], "Inputs": [ { - "OperatorType": "Distinct", - "Collations": [ - "0: utf8mb3_general_ci", - "1: utf8mb3_general_ci", - "2: utf8mb3_general_ci", - "3: utf8mb3_general_ci", - "4: utf8mb3_general_ci", - "5: utf8mb3_general_ci", - "6: utf8mb3_general_ci", - "7", - "8", - "9: utf8mb3_general_ci", - "10: utf8mb3_general_ci", - "11: utf8mb3_general_ci" - ], + "OperatorType": "Concatenate", "Inputs": [ { - "OperatorType": "Concatenate", - "Inputs": [ - { - "OperatorType": "Route", - "Variant": "DBA", - "Keyspace": { - "Name": "main", - "Sharded": false - }, - "FieldQuery": "select `CONSTRAINT_CATALOG`, `CONSTRAINT_SCHEMA`, `CONSTRAINT_NAME`, TABLE_CATALOG, TABLE_SCHEMA, `TABLE_NAME`, `COLUMN_NAME`, ORDINAL_POSITION, POSITION_IN_UNIQUE_CONSTRAINT, REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME from information_schema.key_column_usage as kcu where 1 != 1", - "Query": "select distinct `CONSTRAINT_CATALOG`, `CONSTRAINT_SCHEMA`, `CONSTRAINT_NAME`, TABLE_CATALOG, TABLE_SCHEMA, `TABLE_NAME`, `COLUMN_NAME`, ORDINAL_POSITION, POSITION_IN_UNIQUE_CONSTRAINT, REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME from information_schema.key_column_usage as kcu where kcu.table_schema = :__vtschemaname /* VARCHAR */ and kcu.`table_name` = :kcu_table_name /* VARCHAR */", - "SysTableTableName": "[kcu_table_name:'user_extra']", - "SysTableTableSchema": "['user']", - "Table": "information_schema.key_column_usage" - }, - { - "OperatorType": "Route", - "Variant": "DBA", - "Keyspace": { - "Name": "main", - "Sharded": false - }, - "FieldQuery": "select `CONSTRAINT_CATALOG`, `CONSTRAINT_SCHEMA`, `CONSTRAINT_NAME`, TABLE_CATALOG, TABLE_SCHEMA, `TABLE_NAME`, `COLUMN_NAME`, ORDINAL_POSITION, POSITION_IN_UNIQUE_CONSTRAINT, REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME from information_schema.key_column_usage as kcu where 1 != 1", - "Query": "select distinct `CONSTRAINT_CATALOG`, `CONSTRAINT_SCHEMA`, `CONSTRAINT_NAME`, TABLE_CATALOG, TABLE_SCHEMA, `TABLE_NAME`, `COLUMN_NAME`, ORDINAL_POSITION, POSITION_IN_UNIQUE_CONSTRAINT, REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME from information_schema.key_column_usage as kcu where kcu.table_schema = :__vtschemaname /* VARCHAR */ and kcu.`table_name` = :kcu_table_name1 /* VARCHAR */", - "SysTableTableName": "[kcu_table_name1:'music']", - "SysTableTableSchema": "['user']", - "Table": "information_schema.key_column_usage" - } - ] + "OperatorType": "Route", + "Variant": "DBA", + "Keyspace": { + "Name": "main", + "Sharded": false + }, + "FieldQuery": "select `CONSTRAINT_CATALOG`, `CONSTRAINT_SCHEMA`, `CONSTRAINT_NAME`, TABLE_CATALOG, TABLE_SCHEMA, `TABLE_NAME`, `COLUMN_NAME`, ORDINAL_POSITION, POSITION_IN_UNIQUE_CONSTRAINT, REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME from information_schema.key_column_usage as kcu where 1 != 1", + "Query": "select distinct `CONSTRAINT_CATALOG`, `CONSTRAINT_SCHEMA`, `CONSTRAINT_NAME`, TABLE_CATALOG, TABLE_SCHEMA, `TABLE_NAME`, `COLUMN_NAME`, ORDINAL_POSITION, POSITION_IN_UNIQUE_CONSTRAINT, REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME from information_schema.key_column_usage as kcu where kcu.table_schema = :__vtschemaname /* VARCHAR */ and kcu.`table_name` = :kcu_table_name /* VARCHAR */ and `CONSTRAINT_NAME` = 'primary'", + "SysTableTableName": "[kcu_table_name:'user_extra']", + "SysTableTableSchema": "['user']", + "Table": "information_schema.key_column_usage" + }, + { + "OperatorType": "Route", + "Variant": "DBA", + "Keyspace": { + "Name": "main", + "Sharded": false + }, + "FieldQuery": "select `CONSTRAINT_CATALOG`, `CONSTRAINT_SCHEMA`, `CONSTRAINT_NAME`, TABLE_CATALOG, TABLE_SCHEMA, `TABLE_NAME`, `COLUMN_NAME`, ORDINAL_POSITION, POSITION_IN_UNIQUE_CONSTRAINT, REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME from information_schema.key_column_usage as kcu where 1 != 1", + "Query": "select distinct `CONSTRAINT_CATALOG`, `CONSTRAINT_SCHEMA`, `CONSTRAINT_NAME`, TABLE_CATALOG, TABLE_SCHEMA, `TABLE_NAME`, `COLUMN_NAME`, ORDINAL_POSITION, POSITION_IN_UNIQUE_CONSTRAINT, REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME from information_schema.key_column_usage as kcu where kcu.table_schema = :__vtschemaname /* VARCHAR */ and kcu.`table_name` = :kcu_table_name1 /* VARCHAR */ and `CONSTRAINT_NAME` = 'primary'", + "SysTableTableName": "[kcu_table_name1:'music']", + "SysTableTableSchema": "['user']", + "Table": "information_schema.key_column_usage" } ] }