Add support for drop branch for Iceberg Table

Author: agrawalreetika

presto-docs/src/main/sphinx/connector/iceberg.rst

@@ -1404,6 +1404,8 @@ Alter table operations are supported in the Iceberg connector::
 
      ALTER TABLE iceberg.web.page_views DROP COLUMN location;
 
+     ALTER TABLE iceberg.web.page_views DROP BRANCH 'branch1';
+
 To add a new column as a partition column, identify the transform functions for the column.
 The table is partitioned by the transformed value of the column::
 

presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergAbstractMetadata.java

@@ -183,6 +183,7 @@
 import static com.facebook.presto.iceberg.optimizer.IcebergPlanOptimizer.getEnforcedColumns;
 import static com.facebook.presto.iceberg.util.StatisticsUtil.calculateBaseTableStatistics;
 import static com.facebook.presto.iceberg.util.StatisticsUtil.calculateStatisticsConsideringLayout;
+import static com.facebook.presto.spi.StandardErrorCode.NOT_FOUND;
 import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
 import static com.facebook.presto.spi.statistics.TableStatisticType.ROW_COUNT;
 import static com.google.common.base.Strings.isNullOrEmpty;
@@ -843,6 +844,22 @@ public void rollback()
         // TODO: cleanup open transaction
     }
 
+    @Override
+    public void dropBranch(ConnectorSession session, ConnectorTableHandle tableHandle, String branchName, boolean branchExists)
+    {
+        IcebergTableHandle icebergTableHandle = (IcebergTableHandle) tableHandle;
+        verify(icebergTableHandle.getIcebergTableName().getTableType() == DATA, "only the data table can have branch dropped");
+        Table icebergTable = getIcebergTable(session, icebergTableHandle.getSchemaTableName());
+        if (icebergTable.refs().containsKey(branchName) && icebergTable.refs().get(branchName).isBranch()) {
+            icebergTable.manageSnapshots().removeBranch(branchName).commit();
+        }
+        else {
+            if (!branchExists) {
+                throw new PrestoException(NOT_FOUND, format("Branch %s doesn't exist in table %s", branchName, icebergTableHandle.getSchemaTableName().getTableName()));
+            }
+        }
+    }
+
     @Override
     public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnMetadata column)
     {

presto-iceberg/src/test/java/com/facebook/presto/iceberg/IcebergDistributedTestBase.java

@@ -1817,6 +1817,33 @@ public void testDecimal(boolean decimalVectorReaderEnabled)
         }
     }
 
+    @Test
+    public void testDropBranch()
+    {
+        assertUpdate("CREATE TABLE test_table_branch (id1 BIGINT, id2 BIGINT)");
+        assertUpdate("INSERT INTO test_table_branch VALUES (0, 00), (1, 10)", 2);
+
+        Table icebergTable = loadTable("test_table_branch");
+        icebergTable.manageSnapshots().createBranch("testBranch1").commit();
+
+        assertUpdate("INSERT INTO test_table_branch VALUES (2, 30), (3, 30)", 2);
+        icebergTable.manageSnapshots().createBranch("testBranch2").commit();
+        assertUpdate("INSERT INTO test_table_branch VALUES (4, 40), (5, 50)", 2);
+        assertEquals(icebergTable.refs().size(), 3);
+
+        assertQuery("SELECT count(*) FROM test_table_branch FOR SYSTEM_VERSION AS OF 'testBranch1'", "VALUES 2");
+        assertQuery("SELECT count(*) FROM test_table_branch FOR SYSTEM_VERSION AS OF 'testBranch2'", "VALUES 4");
+        assertQuery("SELECT count(*) FROM test_table_branch FOR SYSTEM_VERSION AS OF 'main'", "VALUES 6");
+
+        assertQuerySucceeds("ALTER TABLE test_table_branch DROP BRANCH 'testBranch1'");
+        icebergTable = loadTable("test_table_branch");
+        assertEquals(icebergTable.refs().size(), 2);
+        assertQueryFails("ALTER TABLE test_table_branch DROP BRANCH 'testBranchNotExist'", "Branch testBranchNotExist doesn't exist in table test_table_branch");
+        assertQuerySucceeds("ALTER TABLE test_table_branch DROP BRANCH IF EXISTS 'testBranch2'");
+        assertQuerySucceeds("ALTER TABLE test_table_branch DROP BRANCH IF EXISTS 'testBranchNotExist'");
+        assertQuerySucceeds("DROP TABLE test_table_branch");
+    }
+
     @Test
     public void testRefsTable()
     {