Repository: vitessio/vitess
Pull request #9703 (2022-02-16 08:46:16)

Push projection to union
## Description

This pull request adds support to push expressions to a union. Previously, queries like:
```sql
select 1
from (
  select id+42 as foo
  from user
  union
  select 1+id as foo
  from unsharded
) as t
```
Would fail because the `1` was not allowed to be pushed down.

## Checklist
- [x] Tests were added or are not required
- [x] Documentation was added or is not required


Patch:
diff --git a/go/vt/vtgate/planbuilder/horizon_planning.go b/go/vt/vtgate/planbuilder/horizon_planning.go
index 1ac8fb25a3..9099a8a8a8 100644
--- a/go/vt/vtgate/planbuilder/horizon_planning.go
+++ b/go/vt/vtgate/planbuilder/horizon_planning.go
@@ -191,7 +191,7 @@ func pushProjection(ctx *plancontext.PlanningContext, expr *sqlparser.AliasedExp
 		expr.Expr = sqlparser.RemoveKeyspaceFromColName(expr.Expr)
 		sel, isSel := node.Select.(*sqlparser.Select)
 		if !isSel {
-			return 0, false, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.BadFieldError, "Unknown column '%s' in 'order clause'", sqlparser.String(expr))
+			return 0, false, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.BadFieldError, "unsupported: pushing projection '%s' on %T", sqlparser.String(expr), node.Select)
 		}
 
 		// if we are trying to push a projection that belongs to a DerivedTable
@@ -390,10 +390,18 @@ func pushProjection(ctx *plancontext.PlanningContext, expr *sqlparser.AliasedExp
 		if err != nil {
 			return 0, false, err
 		}
-		if added {
+		if added && ctx.SemTable.DirectDeps(expr.Expr).NumberOfTables() > 0 {
 			return 0, false, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "pushing projection %v on concatenate should reference an existing column", sqlparser.String(expr))
 		}
-		return offset, false, nil
+		if added {
+			for _, source := range node.sources[1:] {
+				_, _, err := pushProjection(ctx, expr, source, inner, reuseCol, hasAggregation)
+				if err != nil {
+					return 0, false, err
+				}
+			}
+		}
+		return offset, added, nil
 	default:
 		return 0, false, vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "[BUG] push projection does not yet support: %T", node)
 	}


Test patch:
diff --git a/go/vt/vtgate/planbuilder/testdata/union_cases.txt b/go/vt/vtgate/planbuilder/testdata/union_cases.txt
index 77f5c033ec..3ab3ff0e67 100644
--- a/go/vt/vtgate/planbuilder/testdata/union_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/union_cases.txt
@@ -1799,3 +1799,50 @@ Gen4 error: Table 'user' from one of the SELECTs cannot be used in global ORDER
 "select id from user union select 3 order by id"
 "can't do ORDER BY on top of UNION"
 Gen4 plan same as above
+
+"select 1 from (select id+42 as foo from user union select 1+id as foo from unsharded) as t"
+"unsupported: expression on results of a cross-shard subquery"
+{
+  "QueryType": "SELECT",
+  "Original": "select 1 from (select id+42 as foo from user union select 1+id as foo from unsharded) as t",
+  "Instructions": {
+    "OperatorType": "SimpleProjection",
+    "Columns": [
+      1
+    ],
+    "Inputs": [
+      {
+        "OperatorType": "Distinct",
+        "Inputs": [
+          {
+            "OperatorType": "Concatenate",
+            "Inputs": [
+              {
+                "OperatorType": "Route",
+                "Variant": "Scatter",
+                "Keyspace": {
+                  "Name": "user",
+                  "Sharded": true
+                },
+                "FieldQuery": "select id + 42 as foo, 1 from `user` where 1 != 1",
+                "Query": "select distinct id + 42 as foo, 1 from `user`",
+                "Table": "`user`"
+              },
+              {
+                "OperatorType": "Route",
+                "Variant": "Unsharded",
+                "Keyspace": {
+                  "Name": "main",
+                  "Sharded": false
+                },
+                "FieldQuery": "select 1 + id as foo, 1 from unsharded where 1 != 1",
+                "Query": "select distinct 1 + id as foo, 1 from unsharded",
+                "Table": "unsharded"
+              }
+            ]
+          }
+        ]
+      }
+    ]
+  }
+}