Fix iceberg v2 table with equality delete can't update

trinodb · Apr 1, 2024 · 6c4f5dc · 6c4f5dc
1 parent 4c62212
commit 6c4f5dc
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 3 deletions.
diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergUtil.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergUtil.java
@@ -324,9 +324,19 @@ public static IcebergColumnHandle getColumnHandle(NestedField column, TypeManage
 
     public static Schema schemaFromHandles(List<IcebergColumnHandle> columns)
     {
-        List<NestedField> icebergColumns = columns.stream()
-                .map(column -> NestedField.optional(column.getId(), column.getName(), toIcebergType(column.getType(), column.getColumnIdentity())))
-                .collect(toImmutableList());
+        List<NestedField> icebergColumns = new ArrayList<>();
+        for (IcebergColumnHandle column : columns) {
+            NestedField field;
+            // If column is rowIdColumn, we should not create a field with its actual type, because column id will conflict
+            // with a real column. We just use boolean type make a fake.
+            if (column.isUpdateRowIdColumn() || column.isMergeRowIdColumn()) {
+                field = NestedField.optional(column.getId(), column.getName(), toIcebergType(BOOLEAN, column.getColumnIdentity()));
+            }
+            else {
+                field = NestedField.optional(column.getId(), column.getName(), toIcebergType(column.getType(), column.getColumnIdentity()));
+            }
+            icebergColumns.add(field);
+        }
         return new Schema(StructType.of(icebergColumns).asStructType().fields());
     }
 

diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergV2.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergV2.java
@@ -237,6 +237,28 @@ public void testV2TableWithEqualityDeleteWhenColumnIsNested()
         assertQuery("SELECT array_column[1], map_column[1], row_column.x FROM " + tableName, "SELECT 1, 2, 1 FROM nation WHERE regionkey != 1");
     }
 
+    @Test
+    public void testUpdateV2Table()
+    {
+        String tableName = "test_update_v2_table" + randomNameSuffix();
+        assertUpdate("CREATE TABLE " + tableName + " AS SELECT * FROM tpch.tiny.nation", 25);
+        updateTableToV2(tableName);
+        assertUpdate("UPDATE " + tableName + " SET nationkey = 0", 25);
+        assertQuery("SELECT nationkey FROM " + tableName, "SELECT 0 FROM nation");
+    }
+
+    @Test
+    public void testUpdateV2tableWithEqualityDelete()
+            throws Exception
+    {
+        String tableName = "test_update_v2_table" + randomNameSuffix();
+        assertUpdate("CREATE TABLE " + tableName + " AS SELECT * FROM tpch.tiny.nation", 25);
+        Table icebergTable = updateTableToV2(tableName);
+        writeEqualityDeleteToNationTable(icebergTable, Optional.of(icebergTable.spec()), Optional.of(new PartitionData(new Long[]{1L})));
+        assertUpdate("UPDATE " + tableName + " SET nationkey = 0", 20);
+        assertQuery("SELECT nationkey, comment FROM " + tableName, "SELECT 0, comment FROM nation WHERE regionkey != 1");
+    }
+
     @Test
     public void testOptimizingV2TableRemovesEqualityDeletesWhenWholeTableIsScanned()
             throws Exception