Core: Fix move/update/updateDoc/makeRequire/makeOptional fail after r…

…ename schema (#10830)
apache · Feb 10, 2025 · af2c9fb · af2c9fb
1 parent d935460
commit af2c9fb
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 9 deletions.
diff --git a/core/src/main/java/org/apache/iceberg/SchemaUpdate.java b/core/src/main/java/org/apache/iceberg/SchemaUpdate.java
@@ -240,7 +240,7 @@ public UpdateSchema makeColumnOptional(String name) {
   }
 
   private void internalUpdateColumnRequirement(String name, boolean isOptional) {
-    Types.NestedField field = findField(name);
+    Types.NestedField field = findFieldFromUpdateFirst(name);
     Preconditions.checkArgument(field != null, "Cannot update missing column: %s", name);
 
     if ((!isOptional && field.isRequired()) || (isOptional && field.isOptional())) {
@@ -273,7 +273,7 @@ private void internalUpdateColumnRequirement(String name, boolean isOptional) {
 
   @Override
   public UpdateSchema updateColumn(String name, Type.PrimitiveType newType) {
-    Types.NestedField field = findField(name);
+    Types.NestedField field = findFieldFromUpdateFirst(name);
     Preconditions.checkArgument(field != null, "Cannot update missing column: %s", name);
     Preconditions.checkArgument(
         !deletes.contains(field.fieldId()),
@@ -309,7 +309,7 @@ public UpdateSchema updateColumn(String name, Type.PrimitiveType newType) {
 
   @Override
   public UpdateSchema updateColumnDoc(String name, String doc) {
-    Types.NestedField field = findField(name);
+    Types.NestedField field = findFieldFromUpdateFirst(name);
     Preconditions.checkArgument(field != null, "Cannot update missing column: %s", name);
     Preconditions.checkArgument(
         !deletes.contains(field.fieldId()),
@@ -391,13 +391,14 @@ private Integer findForMove(String name) {
     if (addedId != null) {
       return addedId;
     }
+    Types.NestedField field = findFieldFromUpdateFirst(name);
+    return field != null ? field.fieldId() : null;
+  }
 
-    Types.NestedField field = findField(name);
-    if (field != null) {
-      return field.fieldId();
-    }
-
-    return null;
+  private Types.NestedField findFieldFromUpdateFirst(String name) {
+    List<Types.NestedField> updatedFields =
+        updates.values().stream().filter(f -> f.name().equals(name)).collect(Collectors.toList());
+    return !updatedFields.isEmpty() ? updatedFields.get(0) : findField(name);
   }
 
   private void internalMove(String name, Move move) {

diff --git a/core/src/test/java/org/apache/iceberg/TestSchemaUpdate.java b/core/src/test/java/org/apache/iceberg/TestSchemaUpdate.java
@@ -2193,4 +2193,52 @@ public void testMoveDeletedNestedStructFieldToFirst() {
 
     assertThat(actual.asStruct()).isEqualTo(expected.asStruct());
   }
+
+  @Test
+  public void testUpdateDocAfterRename() {
+    Schema schema =
+        new Schema(
+            required(1, "b", Types.IntegerType.get()), required(2, "c", Types.IntegerType.get()));
+
+    Schema actual =
+        new SchemaUpdate(schema, 2).renameColumn("c", "a").updateColumnDoc("a", "doc of a").apply();
+
+    Schema expected =
+        new Schema(
+            required(1, "b", Types.IntegerType.get()),
+            required(2, "a", Types.IntegerType.get(), "doc of a"));
+    assertThat(actual.asStruct()).isEqualTo(expected.asStruct());
+  }
+
+  @Test
+  public void testUpdateAfterRename() {
+    Schema schema =
+        new Schema(
+            required(1, "b", Types.IntegerType.get()), required(2, "c", Types.IntegerType.get()));
+
+    Schema actual =
+        new SchemaUpdate(schema, 2)
+            .renameColumn("c", "a")
+            .updateColumn("a", Types.LongType.get())
+            .apply();
+
+    Schema expected =
+        new Schema(
+            required(1, "b", Types.IntegerType.get()), required(2, "a", Types.LongType.get()));
+    assertThat(actual.asStruct()).isEqualTo(expected.asStruct());
+  }
+
+  @Test
+  public void testMoveAfterRename() {
+    Schema schema =
+        new Schema(
+            required(1, "b", Types.IntegerType.get()), required(2, "c", Types.IntegerType.get()));
+
+    Schema actual = new SchemaUpdate(schema, 2).renameColumn("c", "a").moveBefore("a", "b").apply();
+
+    Schema expected =
+        new Schema(
+            required(2, "a", Types.IntegerType.get()), required(1, "b", Types.IntegerType.get()));
+    assertThat(actual.asStruct()).isEqualTo(expected.asStruct());
+  }
 }