gen4 planner: small cleanup (#11403)

Signed-off-by: Andres Taylor <andres@planetscale.com>

Signed-off-by: Andres Taylor <andres@planetscale.com>

go/vt/vtgate/planbuilder/gen4_planner.go

@@ -222,6 +222,8 @@ func newBuildSelectPlan(
 		return nil, nil, err
 	}
 
+	plan = optimizePlan(plan)
+
 	plan, err = planHorizon(ctx, plan, selStmt, true)
 	if err != nil {
 		return nil, nil, err
@@ -246,6 +248,28 @@ func newBuildSelectPlan(
 	return plan, semTable, nil
 }
 
+// optimizePlan removes unnecessary simpleProjections that have been created while planning
+func optimizePlan(plan logicalPlan) logicalPlan {
+	newPlan, _ := visit(plan, func(plan logicalPlan) (bool, logicalPlan, error) {
+		this, ok := plan.(*simpleProjection)
+		if !ok {
+			return true, plan, nil
+		}
+
+		input, ok := this.input.(*simpleProjection)
+		if !ok {
+			return true, plan, nil
+		}
+
+		for i, col := range this.eSimpleProj.Cols {
+			this.eSimpleProj.Cols[i] = input.eSimpleProj.Cols[col]
+		}
+		this.input = input.input
+		return true, this, nil
+	})
+	return newPlan
+}
+
 func gen4UpdateStmtPlanner(
 	version querypb.ExecuteOptions_PlannerVersion,
 	updStmt *sqlparser.Update,

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

@@ -7288,31 +7288,23 @@ Gen4 plan same as above
         ],
         "Inputs": [
           {
-            "OperatorType": "SimpleProjection",
-            "Columns": [
-              0
-            ],
+            "OperatorType": "Limit",
+            "Count": "INT64(10)",
             "Inputs": [
               {
-                "OperatorType": "Limit",
-                "Count": "INT64(10)",
-                "Inputs": [
-                  {
-                    "OperatorType": "Route",
-                    "Variant": "IN",
-                    "Keyspace": {
-                      "Name": "user",
-                      "Sharded": true
-                    },
-                    "FieldQuery": "select music.id from music where 1 != 1",
-                    "Query": "select music.id from music where music.user_id in ::__vals limit :__upper_limit",
-                    "Table": "music",
-                    "Values": [
-                      "(INT64(5), INT64(6))"
-                    ],
-                    "Vindex": "user_index"
-                  }
-                ]
+                "OperatorType": "Route",
+                "Variant": "IN",
+                "Keyspace": {
+                  "Name": "user",
+                  "Sharded": true
+                },
+                "FieldQuery": "select music.id from music where 1 != 1",
+                "Query": "select music.id from music where music.user_id in ::__vals limit :__upper_limit",
+                "Table": "music",
+                "Values": [
+                  "(INT64(5), INT64(6))"
+                ],
+                "Vindex": "user_index"
               }
             ]
           }
@@ -7422,27 +7414,19 @@ Gen4 plan same as above
         ],
         "Inputs": [
           {
-            "OperatorType": "SimpleProjection",
-            "Columns": [
-              0
-            ],
+            "OperatorType": "Limit",
+            "Count": "INT64(10)",
             "Inputs": [
               {
-                "OperatorType": "Limit",
-                "Count": "INT64(10)",
-                "Inputs": [
-                  {
-                    "OperatorType": "Route",
-                    "Variant": "Scatter",
-                    "Keyspace": {
-                      "Name": "user",
-                      "Sharded": true
-                    },
-                    "FieldQuery": "select music.id from music where 1 != 1",
-                    "Query": "select music.id from music limit :__upper_limit",
-                    "Table": "music"
-                  }
-                ]
+                "OperatorType": "Route",
+                "Variant": "Scatter",
+                "Keyspace": {
+                  "Name": "user",
+                  "Sharded": true
+                },
+                "FieldQuery": "select music.id from music where 1 != 1",
+                "Query": "select music.id from music limit :__upper_limit",
+                "Table": "music"
               }
             ]
           }