Support DROP COLUMN for Hive

trinodb · Jul 11, 2017 · 7629b25 · 7629b25
1 parent d70b505
commit 7629b25
Show file tree

Hide file tree

Showing 10 changed files with 140 additions and 0 deletions.
diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java
@@ -581,6 +581,15 @@ public void renameColumn(ConnectorSession session, ConnectorTableHandle tableHan
         metastore.renameColumn(hiveTableHandle.getSchemaName(), hiveTableHandle.getTableName(), sourceHandle.getName(), target);
     }
 
+    @Override
+    public void dropColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle column)
+    {
+        HiveTableHandle hiveTableHandle = (HiveTableHandle) tableHandle;
+        HiveColumnHandle columnHandle = (HiveColumnHandle) column;
+
+        metastore.dropColumn(hiveTableHandle.getSchemaName(), hiveTableHandle.getTableName(), columnHandle.getName());
+    }
+
     @Override
     public void renameTable(ConnectorSession session, ConnectorTableHandle tableHandle, SchemaTableName newTableName)
     {

diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/metastore/BridgingHiveMetastore.java b/presto-hive/src/main/java/com/facebook/presto/hive/metastore/BridgingHiveMetastore.java
@@ -37,6 +37,7 @@
 import static com.facebook.presto.hive.metastore.MetastoreUtil.toMetastoreApiPartition;
 import static com.facebook.presto.hive.metastore.MetastoreUtil.toMetastoreApiPrivilegeGrantInfo;
 import static com.facebook.presto.hive.metastore.MetastoreUtil.toMetastoreApiTable;
+import static com.facebook.presto.hive.metastore.MetastoreUtil.verifyCanDropColumn;
 import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
 import static java.util.Objects.requireNonNull;
 import static java.util.function.UnaryOperator.identity;
@@ -205,6 +206,20 @@ public void renameColumn(String databaseName, String tableName, String oldColumn
         alterTable(databaseName, tableName, table);
     }
 
+    @Override
+    public void dropColumn(String databaseName, String tableName, String columnName)
+    {
+        verifyCanDropColumn(this, databaseName, tableName, columnName);
+        org.apache.hadoop.hive.metastore.api.Table table = delegate.getTable(databaseName, tableName)
+                .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName)));
+        for (FieldSchema fieldSchema : table.getSd().getCols()) {
+            if (fieldSchema.getName().equals(columnName)) {
+                table.getSd().getCols().remove(fieldSchema);
+            }
+        }
+        alterTable(databaseName, tableName, table);
+    }
+
     private void alterTable(String databaseName, String tableName, org.apache.hadoop.hive.metastore.api.Table table)
     {
         delegate.alterTable(databaseName, tableName, table);

diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/metastore/CachingHiveMetastore.java b/presto-hive/src/main/java/com/facebook/presto/hive/metastore/CachingHiveMetastore.java
@@ -565,6 +565,17 @@ public void renameColumn(String databaseName, String tableName, String oldColumn
         }
     }
 
+    @Override
+    public void dropColumn(String databaseName, String tableName, String columnName)
+    {
+        try {
+            delegate.dropColumn(databaseName, tableName, columnName);
+        }
+        finally {
+            invalidateTable(databaseName, tableName);
+        }
+    }
+
     protected void invalidateTable(String databaseName, String tableName)
     {
         tableCache.invalidate(new HiveTableName(databaseName, tableName));

diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/metastore/ExtendedHiveMetastore.java b/presto-hive/src/main/java/com/facebook/presto/hive/metastore/ExtendedHiveMetastore.java
@@ -59,6 +59,8 @@ public interface ExtendedHiveMetastore
 
     void renameColumn(String databaseName, String tableName, String oldColumnName, String newColumnName);
 
+    void dropColumn(String databaseName, String tableName, String columnName);
+
     Optional<Partition> getPartition(String databaseName, String tableName, List<String> partitionValues);
 
     Optional<List<String>> getPartitionNames(String databaseName, String tableName);

diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/metastore/MetastoreUtil.java b/presto-hive/src/main/java/com/facebook/presto/hive/metastore/MetastoreUtil.java
@@ -19,6 +19,7 @@
 import com.facebook.presto.hive.TableOfflineException;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.SchemaTableName;
+import com.facebook.presto.spi.TableNotFoundException;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import org.apache.hadoop.hive.common.FileUtils;
@@ -57,6 +58,7 @@
 import static com.facebook.presto.hive.HiveErrorCode.HIVE_INVALID_METADATA;
 import static com.facebook.presto.hive.HiveSplitManager.PRESTO_OFFLINE;
 import static com.facebook.presto.hive.metastore.HivePrivilegeInfo.parsePrivilege;
+import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Strings.emptyToNull;
 import static com.google.common.base.Strings.isNullOrEmpty;
@@ -602,4 +604,17 @@ public static void verifyOnline(SchemaTableName tableName, Optional<String> part
             throw new TableOfflineException(tableName, true, prestoOffline);
         }
     }
+
+    public static void verifyCanDropColumn(ExtendedHiveMetastore metastore, String databaseName, String tableName, String columnName)
+    {
+        Table table = metastore.getTable(databaseName, tableName)
+                .orElseThrow(() -> new TableNotFoundException(new SchemaTableName(databaseName, tableName)));
+
+        if (table.getPartitionColumns().stream().anyMatch(column -> column.getName().equals(columnName))) {
+            throw new PrestoException(NOT_SUPPORTED, "Cannot drop partition columns");
+        }
+        if (table.getDataColumns().size() <= 1) {
+            throw new PrestoException(NOT_SUPPORTED, "Cannot drop the only column in a table");
+        }
+    }
 }
diff --git a/...hive/src/main/java/com/facebook/presto/hive/metastore/SemiTransactionalHiveMetastore.java b/...hive/src/main/java/com/facebook/presto/hive/metastore/SemiTransactionalHiveMetastore.java
@@ -349,6 +349,11 @@ public synchronized void renameColumn(String databaseName, String tableName, Str
         setExclusive((delegate, hdfsEnvironment) -> delegate.renameColumn(databaseName, tableName, oldColumnName, newColumnName));
     }
 
+    public synchronized void dropColumn(String databaseName, String tableName, String columnName)
+    {
+        setExclusive((delegate, hdfsEnvironment) -> delegate.dropColumn(databaseName, tableName, columnName));
+    }
+
     public synchronized void finishInsertIntoExistingTable(ConnectorSession session, String databaseName, String tableName, Path currentLocation, List<String> fileNames)
     {
         // Data can only be inserted into partitions and unpartitioned tables. They can never be inserted into a partitioned table.

diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/metastore/file/FileHiveMetastore.java b/presto-hive/src/main/java/com/facebook/presto/hive/metastore/file/FileHiveMetastore.java
@@ -67,6 +67,7 @@
 import static com.facebook.presto.hive.metastore.Database.DEFAULT_DATABASE_NAME;
 import static com.facebook.presto.hive.metastore.HivePrivilegeInfo.HivePrivilege.OWNERSHIP;
 import static com.facebook.presto.hive.metastore.MetastoreUtil.makePartName;
+import static com.facebook.presto.hive.metastore.MetastoreUtil.verifyCanDropColumn;
 import static com.facebook.presto.hive.metastore.PrincipalType.ROLE;
 import static com.facebook.presto.hive.metastore.PrincipalType.USER;
 import static com.facebook.presto.spi.StandardErrorCode.ALREADY_EXISTS;
@@ -427,6 +428,27 @@ public synchronized void renameColumn(String databaseName, String tableName, Str
         });
     }
 
+    @Override
+    public synchronized void dropColumn(String databaseName, String tableName, String columnName)
+    {
+        alterTable(databaseName, tableName, oldTable -> {
+            verifyCanDropColumn(this, databaseName, tableName, columnName);
+            if (!oldTable.getColumn(columnName).isPresent()) {
+                SchemaTableName name = new SchemaTableName(databaseName, tableName);
+                throw new ColumnNotFoundException(name, columnName);
+            }
+
+            ImmutableList.Builder<Column> newDataColumns = ImmutableList.builder();
+            for (Column fieldSchema : oldTable.getDataColumns()) {
+                if (!fieldSchema.getName().equals(columnName)) {
+                    newDataColumns.add(fieldSchema);
+                }
+            }
+
+            return oldTable.withDataColumns(newDataColumns.build());
+        });
+    }
+
     private void alterTable(String databaseName, String tableName, Function<TableMetadata, TableMetadata> alterFunction)
     {
         requireNonNull(databaseName, "databaseName is null");

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
@@ -1861,6 +1861,29 @@ public void testRenameColumn()
         assertUpdate("DROP TABLE test_rename_column");
     }
 
+    @Test
+    public void testDropColumn()
+            throws Exception
+    {
+        @Language("SQL") String createTable = "" +
+                "CREATE TABLE test_drop_column\n" +
+                "WITH (\n" +
+                "  partitioned_by = ARRAY ['orderstatus']\n" +
+                ")\n" +
+                "AS\n" +
+                "SELECT custkey, orderkey, orderstatus FROM orders";
+
+        assertUpdate(createTable, "SELECT count(*) FROM orders");
+        assertQuery("SELECT orderkey, orderstatus FROM test_drop_column", "SELECT orderkey, orderstatus FROM orders");
+
+        assertQueryFails("ALTER TABLE test_drop_column DROP COLUMN orderstatus", "Cannot drop partition columns");
+        assertUpdate("ALTER TABLE test_drop_column DROP COLUMN orderkey");
+        assertQueryFails("ALTER TABLE test_drop_column DROP COLUMN custkey", "Cannot drop the only column in a table");
+        assertQuery("SELECT * FROM test_drop_column", "SELECT custkey, orderstatus FROM orders");
+
+        assertUpdate("DROP TABLE test_drop_column");
+    }
+
     @Test
     public void testAvroTypeValidation()
     {

diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/metastore/TestingHiveMetastore.java b/presto-hive/src/test/java/com/facebook/presto/hive/metastore/TestingHiveMetastore.java
@@ -47,6 +47,7 @@
 import static com.facebook.presto.hive.metastore.Database.DEFAULT_DATABASE_NAME;
 import static com.facebook.presto.hive.metastore.HivePrivilegeInfo.HivePrivilege.OWNERSHIP;
 import static com.facebook.presto.hive.metastore.MetastoreUtil.makePartName;
+import static com.facebook.presto.hive.metastore.MetastoreUtil.verifyCanDropColumn;
 import static com.facebook.presto.hive.metastore.PrincipalType.ROLE;
 import static com.facebook.presto.hive.metastore.PrincipalType.USER;
 import static com.facebook.presto.spi.StandardErrorCode.ALREADY_EXISTS;
@@ -353,6 +354,30 @@ public synchronized void renameColumn(String databaseName, String tableName, Str
         relations.put(name, newTable);
     }
 
+    @Override
+    public synchronized void dropColumn(String databaseName, String tableName, String columnName)
+    {
+        SchemaTableName name = new SchemaTableName(databaseName, tableName);
+        Table oldTable = getRequiredTable(name);
+
+        verifyCanDropColumn(this, databaseName, tableName, columnName);
+        if (!oldTable.getColumn(columnName).isPresent()) {
+            throw new ColumnNotFoundException(name, columnName);
+        }
+
+        ImmutableList.Builder<Column> newDataColumns = ImmutableList.builder();
+        for (Column fieldSchema : oldTable.getDataColumns()) {
+            if (!fieldSchema.getName().equals(columnName)) {
+                newDataColumns.add(fieldSchema);
+            }
+        }
+
+        Table newTable = Table.builder(oldTable)
+                .setDataColumns(newDataColumns.build())
+                .build();
+        relations.put(name, newTable);
+    }
+
     @Override
     public synchronized Optional<List<String>> getAllTables(String databaseName)
     {

diff --git a/presto-product-tests/src/main/java/com/facebook/presto/tests/AlterTableTests.java b/presto-product-tests/src/main/java/com/facebook/presto/tests/AlterTableTests.java
@@ -96,4 +96,17 @@ public void addColumn()
         assertThat(() -> query(format("ALTER TABLE %s ADD COLUMN n_naTioNkEy BIGINT", TABLE_NAME)))
                 .failsWithMessage("Column 'n_naTioNkEy' already exists");
     }
+
+    @Test(groups = {ALTER_TABLE, SMOKE})
+    public void dropColumn()
+    {
+        query(format("CREATE TABLE %s AS SELECT n_nationkey, n_regionkey FROM nation", TABLE_NAME));
+
+        assertThat(query(format("SELECT count(n_nationkey) FROM %s", TABLE_NAME)))
+                .containsExactly(row(25));
+        assertThat(query(format("ALTER TABLE %s DROP COLUMN n_nationkey", TABLE_NAME)))
+                .hasRowsCount(1);
+        assertThat(() -> query(format("ALTER TABLE %s DROP COLUMN n_regionkey", TABLE_NAME)))
+                .failsWithMessage("Cannot drop the only column in a table");
+    }
 }