Improve error message when trying to drop partitioned columns in Iceberg

trinodb · Jan 16, 2023 · 2cc7499 · 2cc7499
1 parent c41c3e1
commit 2cc7499
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java
@@ -131,6 +131,7 @@
 import org.apache.iceberg.expressions.Expressions;
 import org.apache.iceberg.expressions.Term;
 import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.transforms.Transforms;
 import org.apache.iceberg.types.Type;
 import org.apache.iceberg.types.Types;
 import org.apache.iceberg.types.Types.IntegerType;
@@ -1529,6 +1530,11 @@ public void dropColumn(ConnectorSession session, ConnectorTableHandle tableHandl
     {
         IcebergColumnHandle handle = (IcebergColumnHandle) column;
         Table icebergTable = catalog.loadTable(session, ((IcebergTableHandle) tableHandle).getSchemaTableName());
+        boolean isPartitionColumn = icebergTable.spec().fields().stream()
+                .anyMatch(field -> field.sourceId() == handle.getId() && !isVoidTransform(field));
+        if (isPartitionColumn) {
+            throw new TrinoException(NOT_SUPPORTED, "Cannot drop partition field: " + handle.getName());
+        }
         try {
             icebergTable.updateSchema()
                     .deleteColumn(handle.getName())
@@ -1539,6 +1545,11 @@ public void dropColumn(ConnectorSession session, ConnectorTableHandle tableHandl
         }
     }
 
+    private static boolean isVoidTransform(PartitionField field)
+    {
+        return field.transform().equals(Transforms.alwaysNull());
+    }
+
     @Override
     public void renameColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle source, String target)
     {

diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java
@@ -1186,6 +1186,17 @@ public void testSchemaEvolution()
         dropTable("test_schema_evolution_drop_middle");
     }
 
+    @Test
+    public void testDropPartitionColumn()
+    {
+        String tableName = "test_drop_partition_column_" + randomNameSuffix();
+        assertUpdate("CREATE TABLE " + tableName + " (id INTEGER, name VARCHAR, age INTEGER) WITH (partitioning = ARRAY['id', 'truncate(name, 5)', 'void(age)'])");
+        assertQueryFails("ALTER TABLE " + tableName + " DROP COLUMN id", "Cannot drop partition field: id");
+        assertQueryFails("ALTER TABLE " + tableName + " DROP COLUMN name", "Cannot drop partition field: name");
+        assertUpdate("ALTER TABLE " + tableName + " DROP COLUMN age");
+        dropTable(tableName);
+    }
+
     @Test
     public void testShowStatsAfterAddColumn()
     {