[apache#2450] add RenameColumn support for spark-connector

Yangxuhao123 · Mar 12, 2024 · d4c1815 · d4c1815
1 parent a9c6d2b
commit d4c1815
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 0 deletions.
diff --git a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/spark/SparkIT.java b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/spark/SparkIT.java
@@ -298,6 +298,26 @@ void testAlterTableAddAndDeleteColumn() {
     checkTableColumns(tableName, simpleTableColumns, getTableInfo(tableName));
   }
 
+  @Test
+  void testAlterTableRenameColumn() {
+    String tableName = "test_rename_column";
+    dropTableIfExists(tableName);
+    List<SparkColumnInfo> simpleTableColumns = getSimpleTableColumn();
+    createSimpleTable(tableName);
+    checkTableColumns(tableName, simpleTableColumns, getTableInfo(tableName));
+
+    String oldColumnName = "age";
+    String newColumnName = "newAge";
+
+    sql(
+        String.format(
+            "ALTER TABLE %S RENAME COLUMNS %S TO %S", tableName, oldColumnName, newColumnName));
+    ArrayList<SparkColumnInfo> renameColumns = new ArrayList<>(simpleTableColumns);
+    renameColumns.remove(SparkColumnInfo.of("age", DataTypes.IntegerType, null));
+    renameColumns.add(SparkColumnInfo.of("newAge", DataTypes.IntegerType, null));
+    checkTableColumns(tableName, renameColumns, getTableInfo(tableName));
+  }
+
   private void checkTableColumns(
       String tableName, List<SparkColumnInfo> columnInfos, SparkTableInfo tableInfo) {
     SparkTableInfoChecker.create()

diff --git a/...ctor/src/main/java/com/datastrato/gravitino/spark/connector/catalog/GravitinoCatalog.java b/...ctor/src/main/java/com/datastrato/gravitino/spark/connector/catalog/GravitinoCatalog.java
@@ -373,6 +373,10 @@ static com.datastrato.gravitino.rel.TableChange transformTableChange(TableChange
       TableChange.DeleteColumn deleteColumn = (TableChange.DeleteColumn) change;
       return com.datastrato.gravitino.rel.TableChange.deleteColumn(
           deleteColumn.fieldNames(), deleteColumn.ifExists());
+    } else if (change instanceof TableChange.RenameColumn) {
+      TableChange.RenameColumn renameColumn = (TableChange.RenameColumn) change;
+      return com.datastrato.gravitino.rel.TableChange.renameColumn(
+          renameColumn.fieldNames(), renameColumn.newName());
     } else {
       throw new UnsupportedOperationException(
           String.format("Unsupported table change %s", change.getClass().getName()));

diff --git a/.../test/java/com/datastrato/gravitino/spark/connector/catalog/TestTransformTableChange.java b/.../test/java/com/datastrato/gravitino/spark/connector/catalog/TestTransformTableChange.java
@@ -39,6 +39,25 @@ void testTransformRemoveProperty() {
     Assertions.assertEquals("key", gravitinoRemoveProperty.getProperty());
   }
 
+  @Test
+  void testTransformRenameColumn() {
+    String[] oldFiledsName = new String[] {"default_name"};
+    String newFiledName = "new_name";
+
+    TableChange.RenameColumn sparkRenameColumn =
+        (TableChange.RenameColumn) TableChange.renameColumn(oldFiledsName, newFiledName);
+    com.datastrato.gravitino.rel.TableChange gravitinoChange =
+        GravitinoCatalog.transformTableChange(sparkRenameColumn);
+
+    Assertions.assertTrue(
+        gravitinoChange instanceof com.datastrato.gravitino.rel.TableChange.RenameColumn);
+    com.datastrato.gravitino.rel.TableChange.RenameColumn gravitinoRenameColumn =
+        (com.datastrato.gravitino.rel.TableChange.RenameColumn) gravitinoChange;
+
+    Assertions.assertEquals(oldFiledsName, gravitinoRenameCeolumn.getFieldName());
+    Assertions.assertEquals(newFiledName, gravitinoRenameColumn.getNewName());
+  }
+
   @Test
   void testTransformAddColumn() {