API, Core: Add withUpdatedColumnDoc API to Schema

apache · Jan 5, 2024 · 82149f6 · 82149f6
1 parent 580e702
commit 82149f6
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 0 deletions.
diff --git a/api/src/main/java/org/apache/iceberg/Schema.java b/api/src/main/java/org/apache/iceberg/Schema.java
@@ -401,6 +401,37 @@ public String idToAlias(Integer fieldId) {
     return null;
   }
 
+  /**
+   * Returns a schema with an updated field doc for the given field.
+   *
+   * @param name field name
+   * @param doc column doc
+   * @return a Schema with the updated field doc
+   * @throws IllegalArgumentException if no field with the given name is found
+   */
+  public Schema withColumnDoc(String name, String doc) {
+    NestedField fieldToUpdate = findField(name);
+    Preconditions.checkArgument(fieldToUpdate != null, "Field %s does not exist", name);
+    NestedField updatedField =
+        NestedField.of(
+            fieldToUpdate.fieldId(),
+            fieldToUpdate.isOptional(),
+            fieldToUpdate.name(),
+            fieldToUpdate.type(),
+            doc);
+
+    List<NestedField> newFields = Lists.newArrayList();
+    newFields.add(updatedField);
+
+    for (NestedField field : columns()) {
+      if (field.fieldId() != updatedField.fieldId()) {
+        newFields.add(field);
+      }
+    }
+
+    return new Schema(newFields, getAliases(), identifierFieldIds());
+  }
+
   /**
    * Returns an accessor for retrieving the data from {@link StructLike}.
    *

diff --git a/core/src/test/java/org/apache/iceberg/view/ViewCatalogTests.java b/core/src/test/java/org/apache/iceberg/view/ViewCatalogTests.java
@@ -1739,4 +1739,61 @@ public void testSqlForInvalidArguments() {
         .isInstanceOf(IllegalArgumentException.class)
         .hasMessage("Invalid dialect: (empty string)");
   }
+
+  @Test
+  public void testReplaceViewWithUpdatedColumnDoc() {
+    TableIdentifier identifier = TableIdentifier.of("ns", "view");
+
+    if (requiresNamespaceCreate()) {
+      catalog().createNamespace(identifier.namespace());
+    }
+
+    View view =
+        catalog()
+            .buildView(identifier)
+            .withSchema(SCHEMA)
+            .withDefaultNamespace(identifier.namespace())
+            .withDefaultCatalog(catalog().name())
+            .withQuery("spark", "select * from ns.tbl")
+            .create();
+
+    view.replaceVersion()
+        .withDefaultCatalog(view.currentVersion().defaultCatalog())
+        .withDefaultNamespace(identifier.namespace())
+        .withSchema(view.schema().withColumnDoc("data", "some docs for data"))
+        .withQuery("spark", view.sqlFor("spark").sql())
+        .commit();
+
+    assertThat(view.schemas().size()).isEqualTo(2);
+    assertThat(view.schema().findField("data").doc()).isEqualTo("some docs for data");
+  }
+
+  @Test
+  public void testReplaceViewColumnDocNonExistingFieldFails() {
+    TableIdentifier identifier = TableIdentifier.of("ns", "view");
+
+    if (requiresNamespaceCreate()) {
+      catalog().createNamespace(identifier.namespace());
+    }
+
+    View view =
+        catalog()
+            .buildView(identifier)
+            .withSchema(SCHEMA)
+            .withDefaultNamespace(identifier.namespace())
+            .withDefaultCatalog(catalog().name())
+            .withQuery("spark", "select * from ns.tbl")
+            .create();
+
+    assertThatThrownBy(
+            () ->
+                view.replaceVersion()
+                    .withDefaultCatalog(view.currentVersion().defaultCatalog())
+                    .withDefaultNamespace(identifier.namespace())
+                    .withSchema(view.schema().withColumnDoc("non_existing", "non_existing"))
+                    .withQuery("spark", view.sqlFor("spark").sql())
+                    .commit())
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("Field non_existing does not exist");
+  }
 }