wip

go/vt/vtgate/planbuilder/testdata/onecase.json

@@ -1,8 +1,29 @@
 [
   {
     "comment": "Add your test case here for debugging and run go test -run=One.",
-    "query": "",
+    "query": "SELECT music_id FROM music_extra WHERE music_id IN (select * from (select music.id from music where music.user_id = 1234 AND music.foo = 'bar' union select music.id from music where music.user_id = 1234 AND music.foo = 'baz') as subquery)",
     "plan": {
+      "QueryType": "SELECT",
+      "Original": "SELECT music_id FROM music_extra WHERE music_id IN (select music.id from music where music.user_id = 1234 AND music.foo = 'bar' union select music.id from music where music.user_id = 1234 AND music.foo = 'baz')",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select music_id from music_extra where 1 != 1",
+        "Query": "select music_id from music_extra where music_id in (select music.id from music where music.user_id = 1234 and music.foo = 'bar' union select music.id from music where music.user_id = 1234 and music.foo = 'baz')",
+        "Table": "music_extra",
+        "Values": [
+          "1234"
+        ],
+        "Vindex": "user_index"
+      },
+      "TablesUsed": [
+        "user.music",
+        "user.music_extra"
+      ]
     }
   }
 ]

go/vt/vtgate/semantics/analyzer.go

@@ -343,6 +343,7 @@ type originable interface {
 	tableSetFor(t *sqlparser.AliasedTableExpr) TableSet
 	depsForExpr(expr sqlparser.Expr) (direct, recursive TableSet, typ evalengine.Type)
 	collationEnv() *collations.Environment
+	tableInfoFor(ts TableSet) (TableInfo, error)
 }
 
 func (a *analyzer) depsForExpr(expr sqlparser.Expr) (direct, recursive TableSet, typ evalengine.Type) {
@@ -352,6 +353,15 @@ func (a *analyzer) depsForExpr(expr sqlparser.Expr) (direct, recursive TableSet,
 	return
 }
 
+// tableInfoFor returns the table info for the table set. It should contains only single table.
+func (a *analyzer) tableInfoFor(id TableSet) (TableInfo, error) {
+	offset := id.TableOffset()
+	if offset < 0 {
+		return nil, ErrNotSingleTable
+	}
+	return a.tables.Tables[offset], nil
+}
+
 func (a *analyzer) collationEnv() *collations.Environment {
 	return a.typer.collationEnv
 }

go/vt/vtgate/semantics/cte_table.go

@@ -71,6 +71,11 @@ func (cte *CTETable) GetVindexTable() *vindexes.Table {
 	return nil
 }
 
+// getColumnVindexForColumn implements the TableInfo interface
+func (cte *CTETable) getColumnVindexForColumn(org originable, columnName string) (vList []*vindexes.ColumnVindex, err error) {
+	return nil, nil
+}
+
 func (cte *CTETable) IsInfSchema() bool {
 	return false
 }

go/vt/vtgate/semantics/derived_table.go

@@ -31,21 +31,78 @@ type DerivedTable struct {
 	tableName       string
 	ASTNode         *sqlparser.AliasedTableExpr
 	columnNames     []string
-	cols            []sqlparser.Expr
 	tables          TableSet
 	isAuthoritative bool
 
+	selectExprColumns
+}
+
+type selectColumns struct {
+	cols      []sqlparser.Expr
 	recursive []TableSet
 	types     []evalengine.Type
 }
 
-type unionInfo struct {
-	isAuthoritative bool
-	recursive       []TableSet
-	types           []evalengine.Type
-	exprs           sqlparser.SelectExprs
+func (s *selectColumns) getVindexForIndex(org originable, idx int) ([]*vindexes.ColumnVindex, error) {
+	expr := s.cols[idx]
+	colName, isCol := expr.(*sqlparser.ColName)
+	if !isCol {
+		return nil, nil
+	}
+	direct, _, _ := org.depsForExpr(colName)
+	ti, err := org.tableInfoFor(direct)
+	if err != nil {
+		return nil, err
+	}
+	vindexTable := ti.GetVindexTable()
+	var result []*vindexes.ColumnVindex
+	for _, vindex := range vindexTable.ColumnVindexes {
+		if vindex.Columns[0].Equal(colName.Name) {
+			result = append(result, vindex)
+		}
+	}
+	return result, nil
 }
 
+func intersect(a, b []*vindexes.ColumnVindex) []*vindexes.ColumnVindex {
+	var result []*vindexes.ColumnVindex
+	for _, vindex := range a {
+		for _, vindex2 := range b {
+			if vindex == vindex2 {
+				result = append(result, vindex)
+			}
+		}
+	}
+	return result
+}
+
+func (s *unionColumns) getVindexForIndex(org originable, offset int) ([]*vindexes.ColumnVindex, error) {
+	var vlist []*vindexes.ColumnVindex
+	for idx, columns := range s.sc {
+		these, err := columns.getVindexForIndex(org, offset)
+		if err != nil {
+			return nil, err
+		}
+		if idx == 0 {
+			vlist = these
+			continue
+		}
+		vlist = intersect(vlist, these)
+	}
+	return vlist, nil
+}
+
+type unionColumns struct {
+	sc []selectColumns
+}
+
+type selectExprColumns interface {
+	getVindexForIndex(org originable, idx int) ([]*vindexes.ColumnVindex, error)
+}
+
+var _ selectExprColumns = (*selectColumns)(nil)
+var _ selectExprColumns = (*unionColumns)(nil)
+
 var _ TableInfo = (*DerivedTable)(nil)
 
 func createDerivedTableForExpressions(
@@ -154,6 +211,22 @@ func (dt *DerivedTable) GetVindexTable() *vindexes.Table {
 	return nil
 }
 
+// getColumnVindexForColumn implements the TableInfo interface
+func (dt *DerivedTable) getColumnVindexForColumn(org originable, columnName string) (vList []*vindexes.ColumnVindex, err error) {
+	offset := -1
+	for idx, name := range dt.columnNames {
+		if name == columnName {
+			offset = idx
+			break
+		}
+	}
+	if offset == -1 {
+		return nil, nil
+	}
+	return dt.selectExprColumns.getVindexForIndex(org, offset)
+
+}
+
 func (dt *DerivedTable) getColumns(bool) []ColumnInfo {
 	cols := make([]ColumnInfo, 0, len(dt.columnNames))
 	for _, col := range dt.columnNames {

go/vt/vtgate/semantics/real_table.go

@@ -211,6 +211,16 @@ func (r *RealTable) GetVindexTable() *vindexes.Table {
 	return r.Table
 }
 
+// getColumnVindexForColumn implements the TableInfo interface
+func (r *RealTable) getColumnVindexForColumn(org originable, columnName string) (vList []*vindexes.ColumnVindex, err error) {
+	for _, vindex := range r.Table.ColumnVindexes {
+		if len(vindex.Columns) == 1 && vindex.Columns[0].EqualString(columnName) {
+			vList = append(vList, vindex)
+		}
+	}
+	return vList, nil
+}
+
 // GetVindexHint implements the TableInfo interface
 func (r *RealTable) GetVindexHint() *sqlparser.IndexHint {
 	return r.VindexHint

go/vt/vtgate/semantics/semantic_table.go

@@ -64,6 +64,9 @@ type (
 		getExprFor(s string) (sqlparser.Expr, error)
 		getTableSet(org originable) TableSet
 
+		// getColumnVindexForColumn gets the vindex for the given column if available.
+		getColumnVindexForColumn(org originable, columnName string) ([]*vindexes.ColumnVindex, error)
+
 		// GetMirrorRule returns the vschema mirror rule for this TableInfo
 		GetMirrorRule() *vindexes.MirrorRule
 	}
@@ -200,6 +203,37 @@ func (st *SemTable) CopyDependencies(from, to sqlparser.Expr) {
 	}
 }
 
+var _ originable = (*SemTable)(nil)
+
+func (st *SemTable) tableSetFor(t *sqlparser.AliasedTableExpr) TableSet {
+	return st.TableSetFor(t)
+}
+
+func (st *SemTable) depsForExpr(expr sqlparser.Expr) (TableSet, TableSet, evalengine.Type) {
+	typ, _ := st.TypeForExpr(expr)
+	return st.DirectDeps(expr), st.RecursiveDeps(expr), typ
+}
+
+func (st *SemTable) collationEnv() *collations.Environment {
+	return st.collEnv
+}
+
+func (st *SemTable) tableInfoFor(ts TableSet) (TableInfo, error) {
+	return st.TableInfoFor(ts)
+}
+
+func (st *SemTable) GetVindexesForExpr(expr sqlparser.Expr) ([]*vindexes.ColumnVindex, error) {
+	colName, isCol := expr.(*sqlparser.ColName)
+	if !isCol {
+		return nil, nil
+	}
+	ti, err := st.TableInfoForExpr(colName)
+	if err != nil {
+		return nil, err
+	}
+	return ti.getColumnVindexForColumn(st, colName.Name.String())
+}
+
 // GetChildForeignKeysForTargets gets the child foreign keys as a list for all the target tables.
 func (st *SemTable) GetChildForeignKeysForTargets() (fks []vindexes.ChildFKInfo) {
 	for _, ts := range st.Targets.Constituents() {

go/vt/vtgate/semantics/table_collector.go

@@ -47,6 +47,13 @@ type (
 		// cte is a map of CTE definitions that are used in the query
 		cte map[string]*CTE
 	}
+
+	unionInfo struct {
+		isAuthoritative bool
+		recursive       []TableSet
+		types           []evalengine.Type
+		exprs           sqlparser.SelectExprs
+	}
 )
 
 func newEarlyTableCollector(si SchemaInformation, currentDb string) *earlyTableCollector {

go/vt/vtgate/semantics/vindex_table.go

@@ -51,6 +51,11 @@ func (v *VindexTable) GetVindexTable() *vindexes.Table {
 	return v.Table.GetVindexTable()
 }
 
+// getColumnVindexForColumn implements the TableInfo interface
+func (v *VindexTable) getColumnVindexForColumn(org originable, columnName string) (vList []*vindexes.ColumnVindex, err error) {
+	return nil, nil
+}
+
 // Matches implements the TableInfo interface
 func (v *VindexTable) matches(name sqlparser.TableName) bool {
 	return v.Table.matches(name)

go/vt/vtgate/semantics/vtable.go

@@ -104,6 +104,11 @@ func (v *vTableInfo) GetVindexTable() *vindexes.Table {
 	return nil
 }
 
+// getColumnVindexForColumn implements the TableInfo interface
+func (v *vTableInfo) getColumnVindexForColumn(org originable, columnName string) (vList []*vindexes.ColumnVindex, err error) {
+	return nil, nil
+}
+
 func (v *vTableInfo) getColumns(bool) []ColumnInfo {
 	cols := make([]ColumnInfo, 0, len(v.columnNames))
 	for _, col := range v.columnNames {