Repository: prestodb/presto
Issue: Planning bug when creating bucketed table

This query:

``` sql
CREATE TABLE t
WITH (
    bucketed_by=array['k'],
    bucket_count=64) 
AS (SELECT '' a, '' b, 0 k)
```

fails with:

```
java.lang.IllegalArgumentException: Input symbols do not match output symbols
    at com.google.common.base.Preconditions.checkArgument(Preconditions.java:122)
    at com.facebook.presto.sql.planner.plan.ExchangeNode.<init>(ExchangeNode.java:82)
    at com.facebook.presto.sql.planner.plan.ExchangeNode.partitionedExchange(ExchangeNode.java:130)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitTableWriter(AddExchanges.java:636)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitTableWriter(AddExchanges.java:187)
    at com.facebook.presto.sql.planner.plan.TableWriterNode.accept(TableWriterNode.java:128)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.planChild(AddExchanges.java:1231)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitTableFinish(AddExchanges.java:791)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitTableFinish(AddExchanges.java:187)
    at com.facebook.presto.sql.planner.plan.TableFinishNode.accept(TableFinishNode.java:80)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.planChild(AddExchanges.java:1231)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitOutput(AddExchanges.java:236)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitOutput(AddExchanges.java:187)
    at com.facebook.presto.sql.planner.plan.OutputNode.accept(OutputNode.java:81)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges.optimize(AddExchanges.java:149)
    at com.facebook.presto.sql.planner.LogicalPlanner.plan(LogicalPlanner.java:105)
    at com.facebook.presto.execution.SqlQueryExecution.doAnalyzeQuery(SqlQueryExecution.java:281)
    at com.facebook.presto.execution.SqlQueryExecution.analyzeQuery(SqlQueryExecution.java:260)
    at com.facebook.presto.execution.SqlQueryExecution.start(SqlQueryExecution.java:224)
    at com.facebook.presto.execution.QueuedExecution.lambda$start$1(QueuedExecution.java:62)
    at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1142)
    at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:617)
    at java.lang.Thread.run(Thread.java:745)
```


Changed files:
presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java
presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java

diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java
index 11e80aeab0..ad08eb1289 100644
--- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java
+++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java
@@ -550,6 +550,7 @@ public class TestHiveIntegrationSmokeTest
     public void testCreatePartitionedBucketedTableAsFewRows()
             throws Exception
     {
+        // go through all storage formats to make sure the empty buckets are correctly created
         for (TestingHiveStorageFormat storageFormat : getAllTestingHiveStorageFormat()) {
             testCreatePartitionedBucketedTableAsFewRows(storageFormat.getSession(), storageFormat.getFormat());
         }
@@ -614,11 +615,11 @@ public class TestHiveIntegrationSmokeTest
                 "WITH (" +
                 "format = '" + storageFormat + "', " +
                 "partitioned_by = ARRAY[ 'orderstatus' ], " +
-                "bucketed_by = ARRAY[ 'custkey' ], " +
+                "bucketed_by = ARRAY[ 'custkey', 'custkey2' ], " +
                 "bucket_count = 11 " +
                 ") " +
                 "AS " +
-                "SELECT custkey, comment, orderstatus " +
+                "SELECT custkey, custkey AS custkey2, comment, orderstatus " +
                 "FROM tpch.tiny.orders";
 
         assertUpdate(
@@ -650,15 +651,15 @@ public class TestHiveIntegrationSmokeTest
                 "WITH (" +
                 "format = '" + storageFormat + "', " +
                 "partitioned_by = ARRAY[ 'orderstatus' ], " +
-                "bucketed_by = ARRAY[ 'custkey' ], " +
+                "bucketed_by = ARRAY[ 'custkey', 'custkey2' ], " +
                 "bucket_count = 11 " +
                 ") " +
                 "AS " +
-                "SELECT custkey, comment, orderstatus " +
+                "SELECT custkey, custkey AS custkey2, comment, orderstatus " +
                 "FROM tpch.tiny.orders " +
                 "WHERE length(comment) % 2 = 0 " +
                 "UNION ALL " +
-                "SELECT custkey, comment, orderstatus " +
+                "SELECT custkey, custkey AS custkey2, comment, orderstatus " +
                 "FROM tpch.tiny.orders " +
                 "WHERE length(comment) % 2 = 1";
 
@@ -687,22 +688,22 @@ public class TestHiveIntegrationSmokeTest
             assertEquals(columnMetadata.getComment(), annotateColumnComment(Optional.empty(), partitionKey));
         }
 
-        assertEquals(tableMetadata.getMetadata().getProperties().get(BUCKETED_BY_PROPERTY), ImmutableList.of("custkey"));
+        assertEquals(tableMetadata.getMetadata().getProperties().get(BUCKETED_BY_PROPERTY), ImmutableList.of("custkey", "custkey2"));
         assertEquals(tableMetadata.getMetadata().getProperties().get(BUCKET_COUNT_PROPERTY), 11);
 
         List<?> partitions = getPartitions(tableName);
         assertEquals(partitions.size(), 3);
 
-        assertQuery("SELECT * from " + tableName, "SELECT custkey, comment, orderstatus FROM orders");
+        assertQuery("SELECT * from " + tableName, "SELECT custkey, custkey, comment, orderstatus FROM orders");
 
         for (int i = 1; i <= 30; i++) {
             assertQuery(
-                    format("SELECT * from " + tableName + " where custkey = %d", i),
-                    format("SELECT custkey, comment, orderstatus FROM orders where custkey = %d", i));
+                    format("SELECT * from " + tableName + " where custkey = %d and custkey2 = %d", i, i),
+                    format("SELECT custkey, custkey, comment, orderstatus FROM orders where custkey = %d", i));
         }
 
         try {
-            assertUpdate("INSERT INTO " + tableName + " VALUES (1, 'comment', 'O')", 1);
+            assertUpdate("INSERT INTO " + tableName + " VALUES (1, 1, 'comment', 'O')", 1);
             fail("expected failure");
         }
         catch (Exception e) {
@@ -714,6 +715,7 @@ public class TestHiveIntegrationSmokeTest
     public void testInsertPartitionedBucketedTableFewRows()
             throws Exception
     {
+        // go through all storage formats to make sure the empty buckets are correctly created
         for (TestingHiveStorageFormat storageFormat : getAllTestingHiveStorageFormat()) {
             testInsertPartitionedBucketedTableFewRows(storageFormat.getSession(), storageFormat.getFormat());
         }
@@ -801,12 +803,13 @@ public class TestHiveIntegrationSmokeTest
         assertUpdate("" +
                 "CREATE TABLE " + tableName + " (" +
                 "  custkey bigint," +
+                "  custkey2 bigint," +
                 "  comment varchar," +
                 "  orderstatus varchar)" +
                 "WITH (" +
                 "format = '" + storageFormat + "', " +
                 "partitioned_by = ARRAY[ 'orderstatus' ], " +
-                "bucketed_by = ARRAY[ 'custkey' ], " +
+                "bucketed_by = ARRAY[ 'custkey', 'custkey2' ], " +
                 "bucket_count = 11)");
 
         ImmutableList<String> orderStatusList = ImmutableList.of("F", "O", "P");
@@ -817,7 +820,7 @@ public class TestHiveIntegrationSmokeTest
                     getSession().withSystemProperty("task_writer_count", "4"),
                     format(
                             "INSERT INTO " + tableName + " " +
-                                    "SELECT custkey, comment, orderstatus " +
+                                    "SELECT custkey, custkey AS custkey2, comment, orderstatus " +
                                     "FROM tpch.tiny.orders " +
                                     "WHERE orderstatus = '%s'",
                             orderStatus),
@@ -845,12 +848,13 @@ public class TestHiveIntegrationSmokeTest
         assertUpdate("" +
                 "CREATE TABLE " + tableName + " (" +
                 "  custkey bigint," +
+                "  custkey2 bigint," +
                 "  comment varchar," +
                 "  orderstatus varchar)" +
                 "WITH (" +
                 "format = '" + storageFormat + "', " +
                 "partitioned_by = ARRAY[ 'orderstatus' ], " +
-                "bucketed_by = ARRAY[ 'custkey' ], " +
+                "bucketed_by = ARRAY[ 'custkey', 'custkey2' ], " +
                 "bucket_count = 11)");
 
         ImmutableList<String> orderStatusList = ImmutableList.of("F", "O", "P");
@@ -861,11 +865,11 @@ public class TestHiveIntegrationSmokeTest
                     getSession().withSystemProperty("task_writer_count", "4"),
                     format(
                             "INSERT INTO " + tableName + " " +
-                                    "SELECT custkey, comment, orderstatus " +
+                                    "SELECT custkey, custkey AS custkey2, comment, orderstatus " +
                                     "FROM tpch.tiny.orders " +
                                     "WHERE orderstatus = '%s' and length(comment) %% 2 = 0 " +
                                     "UNION ALL " +
-                                    "SELECT custkey, comment, orderstatus " +
+                                    "SELECT custkey, custkey AS custkey2, comment, orderstatus " +
                                     "FROM tpch.tiny.orders " +
                                     "WHERE orderstatus = '%s' and length(comment) %% 2 = 1",
                             orderStatus, orderStatus),
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java
index 8142dbaf18..5111377733 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java
@@ -544,7 +544,7 @@ public class UnaliasSymbolReferences
                     node.getColumnNames(),
                     node.getOutputSymbols(),
                     canonicalize(node.getSampleWeightSymbol()),
-                    node.getPartitioningScheme().map(this::canonicalizePartitionFunctionBinding));
+                    node.getPartitioningScheme().map(partitioningScheme -> canonicalizePartitionFunctionBinding(partitioningScheme, source)));
         }
 
         @Override
@@ -658,11 +658,20 @@ public class UnaliasSymbolReferences
             return builder.build();
         }
 
-        private PartitioningScheme canonicalizePartitionFunctionBinding(PartitioningScheme scheme)
+        private PartitioningScheme canonicalizePartitionFunctionBinding(PartitioningScheme scheme, PlanNode source)
         {
+            Set<Symbol> addedOutputs = new HashSet<>();
+            ImmutableList.Builder<Symbol> outputs = ImmutableList.builder();
+            for (Symbol symbol : source.getOutputSymbols()) {
+                Symbol canonicalOutput = canonicalize(symbol);
+                if (addedOutputs.add(canonicalOutput)) {
+                    outputs.add(canonicalOutput);
+                }
+            }
+
             return new PartitioningScheme(
                     scheme.getPartitioning().translate(this::canonicalize),
-                    canonicalize(scheme.getOutputLayout()),
+                    outputs.build(),
                     canonicalize(scheme.getHashColumn()),
                     scheme.isReplicateNulls(),
                     scheme.getBucketToPartition());