From f4e9965014153f04e491d50c397805d1d99f2d94 Mon Sep 17 00:00:00 2001
From: amogh-jahagirdar <amogh@tabular.io>
Date: Thu, 4 Jan 2024 21:18:03 -0800
Subject: [PATCH] API, Core: Add Schema#withUpdatedDoc and View#updateColumnDoc
 APIs

---
 .../main/java/org/apache/iceberg/Schema.java  | 31 +++++++++++++
 .../java/org/apache/iceberg/view/View.java    | 11 +++++
 .../org/apache/iceberg/view/BaseView.java     | 17 +++++++
 .../apache/iceberg/view/ViewCatalogTests.java | 45 +++++++++++++++++++
 4 files changed, 104 insertions(+)

diff --git a/api/src/main/java/org/apache/iceberg/Schema.java b/api/src/main/java/org/apache/iceberg/Schema.java
index 5e024b7c1c29..7b60b55e77d1 100644
--- 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/api/src/main/java/org/apache/iceberg/view/View.java b/api/src/main/java/org/apache/iceberg/view/View.java
index 779592d03104..2ccc1f4eaaed 100644
--- a/api/src/main/java/org/apache/iceberg/view/View.java
+++ b/api/src/main/java/org/apache/iceberg/view/View.java
@@ -132,4 +132,15 @@ default SQLViewRepresentation sqlFor(String dialect) {
     throw new UnsupportedOperationException(
         "Resolving a sql with a given dialect is not supported");
   }
+
+  /**
+   * Updates the column documentation for the field with the given name
+   *
+   * @param name field name
+   * @param doc doc to update
+   * @return ReplaceViewVersion to replace the current view's version
+   */
+  default ReplaceViewVersion updateColumnDoc(String name, String doc) {
+    throw new UnsupportedOperationException("Updating column documentation is not supported");
+  }
 }
diff --git a/core/src/main/java/org/apache/iceberg/view/BaseView.java b/core/src/main/java/org/apache/iceberg/view/BaseView.java
index d30fc6535098..03e048bccf18 100644
--- a/core/src/main/java/org/apache/iceberg/view/BaseView.java
+++ b/core/src/main/java/org/apache/iceberg/view/BaseView.java
@@ -128,4 +128,21 @@ public SQLViewRepresentation sqlFor(String dialect) {
 
     return closest;
   }
+
+  @Override
+  public ReplaceViewVersion updateColumnDoc(String fieldName, String doc) {
+    ReplaceViewVersion replace =
+        replaceVersion()
+            .withSchema(schema().withColumnDoc(fieldName, doc))
+            .withDefaultCatalog(currentVersion().defaultCatalog())
+            .withDefaultNamespace(currentVersion().defaultNamespace());
+    for (ViewRepresentation representation : currentVersion().representations()) {
+      if (representation instanceof SQLViewRepresentation) {
+        SQLViewRepresentation sqlViewRepresentation = (SQLViewRepresentation) representation;
+        replace.withQuery(sqlViewRepresentation.dialect(), sqlViewRepresentation.sql());
+      }
+    }
+
+    return replace;
+  }
 }
diff --git a/core/src/test/java/org/apache/iceberg/view/ViewCatalogTests.java b/core/src/test/java/org/apache/iceberg/view/ViewCatalogTests.java
index 58727e4588ff..d2e1217d2669 100644
--- 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,49 @@ 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.updateColumnDoc("data", "some docs for data").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.updateColumnDoc("non_existing", "non_existing").commit())
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("Field non_existing does not exist");
+  }
 }