Add layout to ConnectorMetadata beginCreateTable

presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleMetadata.java

@@ -145,7 +145,7 @@ public void renameTable(ConnectorSession session, ConnectorTableHandle tableHand
     @Override
     public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
     {
-        ConnectorOutputTableHandle outputTableHandle = beginCreateTable(session, tableMetadata);
+        ConnectorOutputTableHandle outputTableHandle = beginCreateTable(session, tableMetadata, Optional.empty());
         finishCreateTable(session, outputTableHandle, ImmutableList.of());
     }
 
@@ -170,7 +170,7 @@ public Optional<ConnectorNewTableLayout> getNewTableLayout(ConnectorSession conn
     }
 
     @Override
-    public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Optional<ConnectorNewTableLayout> layout)
     {
         int splitCount = (Integer) tableMetadata.getProperties().get(SPLIT_COUNT_PROPERTY);
         int pagesPerSplit = (Integer) tableMetadata.getProperties().get(PAGES_PER_SPLIT_PROPERTY);

presto-blackhole/src/test/java/com/facebook/presto/plugin/blackhole/TestBlackHoleMetadata.java

@@ -23,6 +23,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
 import static com.facebook.presto.testing.TestingConnectorSession.SESSION;
 import static org.testng.Assert.assertEquals;
@@ -50,7 +51,8 @@ public void tableIsCreatedAfterCommits()
                         schemaTableName,
                         ImmutableList.of(),
                         tableProperties,
-                        null));
+                        null),
+                Optional.empty());
 
         assertThatNoTableIsCreated();
 

presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraMetadata.java

@@ -16,6 +16,7 @@
 import com.facebook.presto.cassandra.util.CassandraCqlUtils;
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.ColumnMetadata;
+import com.facebook.presto.spi.ConnectorNewTableLayout;
 import com.facebook.presto.spi.ConnectorOutputTableHandle;
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.ConnectorTableHandle;
@@ -267,7 +268,7 @@ public void renameTable(ConnectorSession session, ConnectorTableHandle tableHand
     }
 
     @Override
-    public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Optional<ConnectorNewTableLayout> layout)
     {
         checkArgument(!isNullOrEmpty(tableMetadata.getOwner()), "Table owner is null or empty");
 

presto-main/src/main/java/com/facebook/presto/metadata/Metadata.java

@@ -145,7 +145,7 @@ public interface Metadata
     /**
      * Begin the atomic creation of a table with data.
      */
-    OutputTableHandle beginCreateTable(Session session, String catalogName, TableMetadata tableMetadata);
+    OutputTableHandle beginCreateTable(Session session, String catalogName, TableMetadata tableMetadata, Optional<NewTableLayout> layout);
 
     /**
      * Finish a table creation with data after the data is written.

presto-main/src/main/java/com/facebook/presto/metadata/MetadataManager.java

@@ -522,14 +522,14 @@ public Optional<NewTableLayout> getNewTableLayout(Session session, String catalo
     }
 
     @Override
-    public OutputTableHandle beginCreateTable(Session session, String catalogName, TableMetadata tableMetadata)
+    public OutputTableHandle beginCreateTable(Session session, String catalogName, TableMetadata tableMetadata, Optional<NewTableLayout> layout)
     {
         ConnectorEntry entry = connectorsByCatalog.get(catalogName);
         checkArgument(entry != null, "Catalog %s does not exist", catalogName);
         ConnectorMetadata metadata = entry.getMetadataForWrite(session);
         ConnectorTransactionHandle transactionHandle = entry.getTransactionHandle(session);
         ConnectorSession connectorSession = session.toConnectorSession(entry.getCatalog());
-        ConnectorOutputTableHandle handle = metadata.beginCreateTable(connectorSession, tableMetadata.getMetadata());
+        ConnectorOutputTableHandle handle = metadata.beginCreateTable(connectorSession, tableMetadata.getMetadata(), layout.map(NewTableLayout::getLayout));
         return new OutputTableHandle(entry.getConnectorId(), transactionHandle, handle);
     }
 

presto-main/src/main/java/com/facebook/presto/sql/planner/LogicalPlanner.java

@@ -126,7 +126,7 @@ private RelationPlan createTableCreationPlan(Analysis analysis)
         return createTableWriterPlan(
                 analysis,
                 plan,
-                new CreateName(destination.getCatalogName(), tableMetadata),
+                new CreateName(destination.getCatalogName(), tableMetadata, newTableLayout),
                 tableMetadata.getVisibleColumnNames(),
                 newTableLayout);
     }

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

@@ -147,7 +147,7 @@ private TableWriterNode.WriterTarget createWriterTarget(TableWriterNode.WriterTa
             // TODO: begin these operations in pre-execution step, not here
             if (target instanceof TableWriterNode.CreateName) {
                 TableWriterNode.CreateName create = (TableWriterNode.CreateName) target;
-                return new TableWriterNode.CreateHandle(metadata.beginCreateTable(session, create.getCatalog(), create.getTableMetadata()));
+                return new TableWriterNode.CreateHandle(metadata.beginCreateTable(session, create.getCatalog(), create.getTableMetadata(), create.getLayout()));
             }
             if (target instanceof TableWriterNode.InsertReference) {
                 TableWriterNode.InsertReference insert = (TableWriterNode.InsertReference) target;

presto-main/src/main/java/com/facebook/presto/sql/planner/plan/TableWriterNode.java

@@ -14,6 +14,7 @@
 package com.facebook.presto.sql.planner.plan;
 
 import com.facebook.presto.metadata.InsertTableHandle;
+import com.facebook.presto.metadata.NewTableLayout;
 import com.facebook.presto.metadata.OutputTableHandle;
 import com.facebook.presto.metadata.TableHandle;
 import com.facebook.presto.metadata.TableMetadata;
@@ -145,11 +146,13 @@ public static class CreateName
     {
         private final String catalog;
         private final TableMetadata tableMetadata;
+        private final Optional<NewTableLayout> layout;
 
-        public CreateName(String catalog, TableMetadata tableMetadata)
+        public CreateName(String catalog, TableMetadata tableMetadata, Optional<NewTableLayout> layout)
         {
             this.catalog = requireNonNull(catalog, "catalog is null");
             this.tableMetadata = requireNonNull(tableMetadata, "tableMetadata is null");
+            this.layout = requireNonNull(layout, "layout is null");
         }
 
         public String getCatalog()
@@ -162,6 +165,11 @@ public TableMetadata getTableMetadata()
             return tableMetadata;
         }
 
+        public Optional<NewTableLayout> getLayout()
+        {
+            return layout;
+        }
+
         @Override
         public String toString()
         {

presto-main/src/main/java/com/facebook/presto/transaction/LegacyConnectorMetadata.java

@@ -161,7 +161,7 @@ public final Optional<ConnectorNewTableLayout> getNewTableLayout(ConnectorSessio
     }
 
     @Override
-    public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Optional<ConnectorNewTableLayout> layout)
     {
         checkState(rollbackAction.get() == null, "Cannot begin a new write while in an existing one");
         ConnectorOutputTableHandle outputTableHandle = metadata.beginCreateTable(session, tableMetadata);

presto-raptor/src/main/java/com/facebook/presto/raptor/RaptorMetadata.java

@@ -24,6 +24,7 @@
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.ColumnMetadata;
 import com.facebook.presto.spi.ConnectorInsertTableHandle;
+import com.facebook.presto.spi.ConnectorNewTableLayout;
 import com.facebook.presto.spi.ConnectorOutputTableHandle;
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.ConnectorTableHandle;
@@ -256,7 +257,7 @@ public ConnectorTableLayout getTableLayout(ConnectorSession session, ConnectorTa
     @Override
     public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
     {
-        finishCreateTable(session, beginCreateTable(session, tableMetadata), ImmutableList.of());
+        finishCreateTable(session, beginCreateTable(session, tableMetadata, Optional.empty()), ImmutableList.of());
     }
 
     @Override
@@ -303,7 +304,7 @@ public void renameColumn(ConnectorSession session, ConnectorTableHandle tableHan
     }
 
     @Override
-    public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Optional<ConnectorNewTableLayout> layout)
     {
         ImmutableList.Builder<RaptorColumnHandle> columnHandles = ImmutableList.builder();
         ImmutableList.Builder<Type> columnTypes = ImmutableList.builder();

presto-raptor/src/test/java/com/facebook/presto/raptor/metadata/TestRaptorMetadata.java

@@ -50,6 +50,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.OptionalLong;
 import java.util.stream.Collectors;
 
@@ -394,7 +395,7 @@ public void testTransactionTableWrite()
     {
         // start table creation
         long transactionId = 1;
-        ConnectorOutputTableHandle outputHandle = metadata.beginCreateTable(SESSION, getOrdersTable());
+        ConnectorOutputTableHandle outputHandle = metadata.beginCreateTable(SESSION, getOrdersTable(), Optional.empty());
 
         // transaction is in progress
         assertTrue(transactionExists(transactionId));
@@ -479,7 +480,7 @@ public void testTransactionAbort()
     {
         // start table creation
         long transactionId = 1;
-        ConnectorOutputTableHandle outputHandle = metadata.beginCreateTable(SESSION, getOrdersTable());
+        ConnectorOutputTableHandle outputHandle = metadata.beginCreateTable(SESSION, getOrdersTable(), Optional.empty());
 
         // transaction is in progress
         assertTrue(transactionExists(transactionId));

presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorMetadata.java

@@ -172,7 +172,7 @@ default Optional<ConnectorNewTableLayout> getNewTableLayout(ConnectorSession ses
     /**
      * Begin the atomic creation of a table with data.
      */
-    default ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    default ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Optional<ConnectorNewTableLayout> layout)
     {
         throw new PrestoException(NOT_SUPPORTED, "This connector does not support creating tables with data");
     }