From 035f136caa94b5d0aa5e53adbfe6f2b270045874 Mon Sep 17 00:00:00 2001
From: wangd <mingwbd@gmail.com>
Date: Sat, 4 Nov 2023 07:43:20 +0800
Subject: [PATCH] Disallow dropping partition column in Iceberg table
 explicitly

---
 .../presto/iceberg/IcebergAbstractMetadata.java | 11 +++++++++++
 .../iceberg/IcebergDistributedTestBase.java     | 17 +++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergAbstractMetadata.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergAbstractMetadata.java
index 78516400a994..85f303e51096 100644
--- a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergAbstractMetadata.java
+++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergAbstractMetadata.java
@@ -378,6 +378,17 @@ public void dropColumn(ConnectorSession session, ConnectorTableHandle tableHandl
         IcebergTableHandle icebergTableHandle = (IcebergTableHandle) tableHandle;
         IcebergColumnHandle handle = (IcebergColumnHandle) column;
         Table icebergTable = getIcebergTable(session, icebergTableHandle.getSchemaTableName());
+
+        // Currently drop partition column used in any partition specs of a table would introduce some problems in Iceberg.
+        // So we explicitly disallow dropping partition columns until Iceberg fix this problem.
+        // See https://github.com/apache/iceberg/issues/4563
+        boolean shouldNotDropPartitionColumn = icebergTable.specs().values().stream()
+                .flatMap(partitionSpec -> partitionSpec.fields().stream())
+                .anyMatch(field -> field.sourceId() == handle.getId());
+        if (shouldNotDropPartitionColumn) {
+            throw new PrestoException(NOT_SUPPORTED, "This connector does not support dropping columns which exist in any of the table's partition specs");
+        }
+
         icebergTable.updateSchema().deleteColumn(handle.getName()).commit();
     }
 
diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/IcebergDistributedTestBase.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/IcebergDistributedTestBase.java
index b39bf002bf56..84d02f4aacee 100644
--- a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/IcebergDistributedTestBase.java
+++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/IcebergDistributedTestBase.java
@@ -217,6 +217,23 @@ public void testAddPartitionColumn()
         assertQuerySucceeds("DROP TABLE add_partition_column");
     }
 
+    @Test
+    public void testDropPartitionColumn()
+    {
+        assertQuerySucceeds("create table test_drop_partition_column(a int, b varchar) with (partitioning = ARRAY['a'])");
+        assertQuerySucceeds("insert into test_drop_partition_column values(1, '1001'), (2, '1002'), (3, '1003')");
+        String errorMessage = "This connector does not support dropping columns which exist in any of the table's partition specs";
+        assertQueryFails("alter table test_drop_partition_column drop column a", errorMessage);
+        assertQuerySucceeds("DROP TABLE test_drop_partition_column");
+
+        assertQuerySucceeds("create table test_drop_partition_column(a int)");
+        assertQuerySucceeds("insert into test_drop_partition_column values 1, 2, 3");
+        assertQuerySucceeds("alter table test_drop_partition_column add column b varchar with (partitioning = 'identity')");
+        assertQuerySucceeds("insert into test_drop_partition_column values(4, '1004'), (5, '1005')");
+        assertQueryFails("alter table test_drop_partition_column drop column b", errorMessage);
+        assertQuerySucceeds("DROP TABLE test_drop_partition_column");
+    }
+
     @Test
     public void testTruncate()
     {