Simplify union node with empty input

presto-hive/src/test/java/com/facebook/presto/hive/TestHiveMaterializedViewLogicalPlanner.java

@@ -530,10 +530,10 @@ public void testMaterializedViewForUnionAllWithOneSideMaterialized()
             MaterializedResult baseTable = computeActual(baseQuery);
             assertEquals(viewTable, baseTable);
 
-            assertPlan(getSession(), viewQuery, anyTree(values("ds", "orderkey"), anyTree(
+            assertPlan(getSession(), viewQuery, anyTree(
                     constrainedTableScan(table2,
                             ImmutableMap.of("ds", multipleValues(createVarcharType(10), utf8Slices("2019-01-02")))),
-                    constrainedTableScan(view, ImmutableMap.of()))));
+                    constrainedTableScan(view, ImmutableMap.of())));
         }
         finally {
             queryRunner.execute("DROP MATERIALIZED VIEW IF EXISTS " + view);

presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/SimplifyPlanWithEmptyInput.java

@@ -43,6 +43,7 @@
 import com.google.common.collect.ImmutableList;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.stream.IntStream;
 
@@ -71,7 +72,7 @@
  * <pre>
  *     - Empty Values
  * </pre>
- *
+ * <p>
  * For outer join: replace with empty values if outer side is empty and project node if inner side is empty
  * For example:
  * <pre>
@@ -87,16 +88,17 @@
  *          assignments := NULL if output not in Scan, otherwise identity projection
  *          - Scan
  * </pre>
- *
+ * <p>
  * For aggregation: if it has default output for empty input, stop and do not simplify, otherwise convert to empty values node
  * <pre>
  *     - Aggregation
  *          count() without group by
  *          - Empty Values
  * </pre>
  * No change for this query plan
- *
- * For Union node: if it has only one non-empty input, convert to a project node. If all inputs are empty, convert to empty values node
+ * <p>
+ * For Union node: if it has only one non-empty input, convert to a project node. If all inputs are empty, convert to empty values node. If more than one input is non-empty,
+ * remove the empty inputs and keep the union node and the non-empty inputs.
  */
 
 public class SimplifyPlanWithEmptyInput
@@ -223,6 +225,13 @@ else if (nonEmptyChildIndex.size() == 1) {
                 builder.putAll(node.getVariableMapping().entrySet().stream().collect(toImmutableMap(entry -> entry.getKey(), entry -> entry.getValue().get(index))));
                 return new ProjectNode(node.getSourceLocation(), idAllocator.getNextId(), rewrittenChildren.get(index), builder.build(), LOCAL);
             }
+            else if (nonEmptyChildIndex.size() < node.getSources().size()) {
+                this.planChanged = true;
+                List<PlanNode> nonEmptyInput = nonEmptyChildIndex.stream().map(x -> node.getSources().get(x)).collect(toImmutableList());
+                Map<VariableReferenceExpression, List<VariableReferenceExpression>> newOutputToInputs = node.getVariableMapping().entrySet().stream()
+                        .collect(toImmutableMap(Map.Entry::getKey, entry -> nonEmptyChildIndex.stream().map(idx -> entry.getValue().get(idx)).collect(toImmutableList())));
+                return new UnionNode(node.getSourceLocation(), idAllocator.getNextId(), nonEmptyInput, node.getOutputVariables(), newOutputToInputs);
+            }
             return node.replaceChildren(rewrittenChildren);
         }
 

presto-main/src/test/java/com/facebook/presto/sql/planner/optimizations/TestSimplifyPlanWithEmptyInput.java

@@ -21,6 +21,7 @@
 
 import static com.facebook.presto.SystemSessionProperties.SIMPLIFY_PLAN_WITH_EMPTY_INPUT;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.aggregation;
+import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.exchange;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.expression;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.filter;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.functionCall;
@@ -137,6 +138,18 @@ public void testUnionWithEmptyInput()
                         tableScan("lineitem", ImmutableMap.of("partkey", "partkey", "orderkey", "orderkey"))));
     }
 
+    @Test
+    public void testUnionMultipleNonEmptyInput()
+    {
+        assertPlan("select orderkey, partkey from lineitem union all select orderkey, custkey as partkey from orders where false union all select custkey, nationkey from customer",
+                enableOptimization(),
+                output(
+                        ImmutableList.of("orderkey", "partkey"),
+                        exchange(
+                                tableScan("lineitem", ImmutableMap.of("partkey", "partkey", "orderkey", "orderkey")),
+                                tableScan("customer", ImmutableMap.of("nationkey", "nationkey", "custkey", "custkey")))));
+    }
+
     @Test
     public void testSemiJoinEmptyFilterSource()
     {