Merge pull request #9521 from planetscale/gen4-outputcol-fix

Fix: push output columns on lhs than on jointree, handle derived table table push output column correctly

go/vt/vtgate/planbuilder/derivedtree.go

@@ -50,28 +50,41 @@ func (d *derivedTree) clone() queryTree {
 	return &other
 }
 
-func (d *derivedTree) pushOutputColumns(names []*sqlparser.ColName, semTable *semantics.SemTable) (offsets []int, err error) {
+func (d *derivedTree) pushOutputColumns(columns []*sqlparser.ColName, semTable *semantics.SemTable) ([]int, error) {
 	var noQualifierNames []*sqlparser.ColName
-	if len(names) == 0 {
-		return
+	var offsets []int
+	if len(columns) == 0 {
+		return nil, nil
 	}
-	for _, name := range names {
-		i, err := d.findOutputColumn(name)
+	for _, col := range columns {
+		i, err := d.findOutputColumn(col)
 		if err != nil {
 			return nil, err
 		}
+		var pos int
+		d.columnsOffset, pos = addToIntSlice(d.columnsOffset, i)
+		offsets = append(offsets, pos)
+		// skip adding to columns as it exists already.
 		if i > -1 {
-			d.columnsOffset = append(d.columnsOffset, i)
 			continue
 		}
-		d.columnsOffset = append(d.columnsOffset, i)
-		d.columns = append(d.columns, name)
-		noQualifierNames = append(noQualifierNames, sqlparser.NewColName(name.Name.String()))
+		d.columns = append(d.columns, col)
+		noQualifierNames = append(noQualifierNames, sqlparser.NewColName(col.Name.String()))
 	}
 	if len(noQualifierNames) > 0 {
 		_, _ = d.inner.pushOutputColumns(noQualifierNames, semTable)
 	}
-	return d.columnsOffset, nil
+	return offsets, nil
+}
+
+func addToIntSlice(columnOffset []int, valToAdd int) ([]int, int) {
+	for idx, val := range columnOffset {
+		if val == valToAdd {
+			return columnOffset, idx
+		}
+	}
+	columnOffset = append(columnOffset, valToAdd)
+	return columnOffset, len(columnOffset) - 1
 }
 
 func (d *derivedTree) pushPredicate(ctx *planningContext, expr sqlparser.Expr) error {

go/vt/vtgate/planbuilder/route_planning.go

@@ -564,7 +564,7 @@ func pushJoinPredicateOnJoin(ctx *planningContext, exprs []sqlparser.Expr, node
 	}
 
 	if lhsColumns != nil && lhsVarsName != nil {
-		idxs, err := node.pushOutputColumns(lhsColumns, ctx.semTable)
+		idxs, err := node.lhs.pushOutputColumns(lhsColumns, ctx.semTable)
 		if err != nil {
 			return nil, err
 		}

go/vt/vtgate/planbuilder/testdata/filter_cases.txt

@@ -904,7 +904,7 @@ Gen4 plan same as above
     "Variant": "Join",
     "JoinColumnIndexes": "1",
     "JoinVars": {
-      "user_col": 1
+      "user_col": 0
     },
     "Predicate": "`user`.col = user_extra.col and user_extra.user_id = `user`.col",
     "TableName": "`user`_user_extra",

go/vt/vtgate/planbuilder/testdata/tpch_cases.txt

@@ -449,10 +449,10 @@ Gen4 error: unsupported: cross-shard correlated subquery
         "Variant": "Join",
         "JoinColumnIndexes": "-5",
         "JoinVars": {
-          "l_partkey": 10,
-          "l_quantity": 12,
-          "l_shipinstruct": 14,
-          "l_shipmode": 13
+          "l_partkey": 0,
+          "l_quantity": 1,
+          "l_shipinstruct": 3,
+          "l_shipmode": 2
         },
         "Predicate": "p_partkey = l_partkey and p_brand = 'Brand#12' and p_container in ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG') and l_quantity \u003e= 1 and l_quantity \u003c= 1 + 10 and p_size between 1 and 5 and l_shipmode in ('AIR', 'AIR REG') and l_shipinstruct = 'DELIVER IN PERSON' or p_partkey = l_partkey and p_brand = 'Brand#23' and p_container in ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK') and l_quantity \u003e= 10 and l_quantity \u003c= 10 + 10 and p_size between 1 and 10 and l_shipmode in ('AIR', 'AIR REG') and l_shipinstruct = 'DELIVER IN PERSON' or p_partkey = l_partkey and p_brand = 'Brand#34' and p_container in ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG') and l_quantity \u003e= 20 and l_quantity \u003c= 20 + 10 and p_size between 1 and 15 and l_shipmode in ('AIR', 'AIR REG') and l_shipinstruct = 'DELIVER IN PERSON'",
         "TableName": "lineitem_part",