Pass through predicates when trying to merge apply join.

Signed-off-by: Arthur Schreiber <[email protected]>

Author: arthurschreiber

go/vt/vtgate/planbuilder/operators/query_planning.go

@@ -21,6 +21,7 @@ import (
 	"io"
 	"strconv"
 
+	"vitess.io/vitess/go/slice"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/predicates"
 
 	"vitess.io/vitess/go/vt/sqlparser"
@@ -123,7 +124,11 @@ func runRewriters(ctx *plancontext.PlanningContext, root Operator) Operator {
 }
 
 func tryMergeApplyJoin(in *ApplyJoin, ctx *plancontext.PlanningContext) (_ Operator, res *ApplyResult) {
-	jm := newJoinMerge(nil, in.JoinType)
+	preds := slice.Map(in.JoinPredicates.columns, func(col applyJoinColumn) sqlparser.Expr {
+		return col.Original
+	})
+
+	jm := newJoinMerge(preds, in.JoinType)
 	r := jm.mergeJoinInputs(ctx, in.LHS, in.RHS)
 	if r == nil {
 		return in, NoRewrite

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

@@ -5450,6 +5450,33 @@
       ]
     }
   },
+  {
+    "comment": "Joining with a subquery that uses an aggregate column and an `EqualUnique` route with additional subquery can be merged together",
+    "query": "SELECT music.id FROM music INNER JOIN (SELECT MAX(id) as maxt FROM music WHERE music.user_id = 5 AND music.foobar NOT IN (SELECT foobar FROM music_extra WHERE music_extra.user_id = 5)) other ON other.maxt = music.id",
+    "plan": {
+      "Type": "Passthrough",
+      "QueryType": "SELECT",
+      "Original": "SELECT music.id FROM music INNER JOIN (SELECT MAX(id) as maxt FROM music WHERE music.user_id = 5 AND music.foobar NOT IN (SELECT foobar FROM music_extra WHERE music_extra.user_id = 5)) other ON other.maxt = music.id",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select music.id from (select max(id) as maxt from music where 1 != 1) as other, music where 1 != 1",
+        "Query": "select music.id from (select max(id) as maxt from music where music.user_id = 5 and music.foobar not in (select foobar from music_extra where music_extra.user_id = 5)) as other, music where other.maxt = music.id",
+        "Table": "music",
+        "Values": [
+          "5"
+        ],
+        "Vindex": "user_index"
+      },
+      "TablesUsed": [
+        "user.music"
+      ]
+    }
+  },
   {
     "comment": "Joining with a subquery that uses an `EqualUnique` route can be merged",
     "query": "SELECT music.id FROM music INNER JOIN (SELECT id FROM music WHERE music.user_id = 5) other ON other.id = music.id",