[WIP] Rescale bucket

Author: Liebing

fluss-client/src/main/java/org/apache/fluss/client/admin/Admin.java

@@ -234,6 +234,25 @@ CompletableFuture<Void> createTable(
      */
     CompletableFuture<List<String>> listTables(String databaseName);
 
+    /**
+     * Alter a table.
+     *
+     * <p>The following exceptions can be anticipated when calling {@code get()} on returned future.
+     *
+     * <ul>
+     *   <li>{@link DatabaseNotExistException} when the database does not exist.
+     *   <li>{@link TableNotExistException} when the table does not exist, if ignoreIfNotExists is
+     *       false.
+     * </ul>
+     *
+     * @param tablePath The table path of the table.
+     * @param tableDescriptor The table descriptor.
+     * @param ignoreIfNotExists if it is true, do nothing if table does not exist. If false, throw a
+     *     TableNotExistException.
+     */
+    CompletableFuture<Void> alterTable(
+            TablePath tablePath, TableDescriptor tableDescriptor, boolean ignoreIfNotExists);
+
     /**
      * List all partitions in the given table in fluss cluster asynchronously.
      *

fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java

@@ -41,6 +41,7 @@
 import org.apache.fluss.rpc.gateway.AdminGateway;
 import org.apache.fluss.rpc.gateway.AdminReadOnlyGateway;
 import org.apache.fluss.rpc.gateway.TabletServerGateway;
+import org.apache.fluss.rpc.messages.AlterTableRequest;
 import org.apache.fluss.rpc.messages.CreateAclsRequest;
 import org.apache.fluss.rpc.messages.CreateDatabaseRequest;
 import org.apache.fluss.rpc.messages.CreateTableRequest;
@@ -235,6 +236,19 @@ public CompletableFuture<Void> createTable(
         return gateway.createTable(request).thenApply(r -> null);
     }
 
+    @Override
+    public CompletableFuture<Void> alterTable(
+            TablePath tablePath, TableDescriptor tableDescriptor, boolean ignoreIfNotExists) {
+        tablePath.validate();
+        AlterTableRequest request = new AlterTableRequest();
+        request.setTableJson(tableDescriptor.toJsonBytes())
+                .setIgnoreIfNotExists(ignoreIfNotExists)
+                .setTablePath()
+                .setDatabaseName(tablePath.getDatabaseName())
+                .setTableName(tablePath.getTableName());
+        return gateway.alterTable(request).thenApply(r -> null);
+    }
+
     @Override
     public CompletableFuture<TableInfo> getTableInfo(TablePath tablePath) {
         GetTableInfoRequest request = new GetTableInfoRequest();

fluss-client/src/test/java/org/apache/fluss/client/admin/ClientToServerITCaseBase.java

@@ -83,6 +83,9 @@ public abstract class ClientToServerITCaseBase {
     @BeforeEach
     protected void setup() throws Exception {
         clientConf = FLUSS_CLUSTER_EXTENSION.getClientConfig();
+        clientConf.set(
+                ConfigOptions.CLIENT_WRITER_BUCKET_NO_KEY_ASSIGNER,
+                ConfigOptions.NoKeyAssigner.ROUND_ROBIN);
         conn = ConnectionFactory.createConnection(clientConf);
         admin = conn.getAdmin();
     }
@@ -235,6 +238,13 @@ public static void waitAllReplicasReady(long tableId, int expectBucketCount) {
         }
     }
 
+    public static void waitAllReplicasReady(long tableId, long partitionId, int expectBucketCount) {
+        for (int i = 0; i < expectBucketCount; i++) {
+            FLUSS_CLUSTER_EXTENSION.waitUntilAllReplicaReady(
+                    new TableBucket(tableId, partitionId, i));
+        }
+    }
+
     protected static void verifyRows(
             RowType rowType,
             Map<Long, List<InternalRow>> actualRows,

fluss-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java

@@ -60,6 +60,7 @@
 import org.apache.fluss.metadata.TablePath;
 import org.apache.fluss.server.kv.snapshot.CompletedSnapshot;
 import org.apache.fluss.server.kv.snapshot.KvSnapshotHandle;
+import org.apache.fluss.server.zk.data.TableAssignment;
 import org.apache.fluss.types.DataTypes;
 
 import org.junit.jupiter.api.BeforeEach;
@@ -197,6 +198,117 @@ void testGetTableInfoAndSchema() throws Exception {
                 .isBetween(timestampBeforeCreate, timestampAfterCreate);
     }
 
+    @Test
+    void testAlterTableBucket() throws Exception {
+        // create table
+        TablePath tablePath = TablePath.of("test_db", "alter_table_bucket");
+        admin.createTable(tablePath, DEFAULT_TABLE_DESCRIPTOR, false).get();
+
+        TableInfo tableInfo = admin.getTableInfo(tablePath).get();
+
+        TableAssignment tableAssignment =
+                FLUSS_CLUSTER_EXTENSION
+                        .getZooKeeperClient()
+                        .getTableAssignment(tableInfo.getTableId())
+                        .get();
+        System.out.println(tableAssignment);
+
+        TableDescriptor existingTableDescriptor = tableInfo.toTableDescriptor();
+
+        TableDescriptor newTableDescriptor =
+                TableDescriptor.builder()
+                        .schema(existingTableDescriptor.getSchema())
+                        .comment(existingTableDescriptor.getComment().orElse("test table"))
+                        .partitionedBy(existingTableDescriptor.getPartitionKeys())
+                        .distributedBy(
+                                existingTableDescriptor
+                                                .getTableDistribution()
+                                                .get()
+                                                .getBucketCount()
+                                                .get()
+                                        + 1,
+                                existingTableDescriptor.getBucketKeys())
+                        .properties(existingTableDescriptor.getProperties())
+                        .customProperties(existingTableDescriptor.getCustomProperties())
+                        .build();
+        // alter table
+        admin.alterTable(tablePath, newTableDescriptor, false).get();
+
+        TableInfo alteredTableInfo = admin.getTableInfo(tablePath).get();
+        TableDescriptor alteredTableDescriptor = alteredTableInfo.toTableDescriptor();
+        // assertThat(alteredTableDescriptor).isEqualTo(newTableDescriptor);
+
+        TableAssignment tableAssignment1 =
+                FLUSS_CLUSTER_EXTENSION
+                        .getZooKeeperClient()
+                        .getTableAssignment(tableInfo.getTableId())
+                        .get();
+        System.out.println(tableAssignment1);
+    }
+
+    @Test
+    void testAlterTable() throws Exception {
+        // create table
+        TablePath tablePath = TablePath.of("test_db", "alter_table_1");
+        admin.createTable(tablePath, DEFAULT_TABLE_DESCRIPTOR, false).get();
+
+        TableInfo tableInfo = admin.getTableInfo(tablePath).get();
+
+        TableDescriptor existingTableDescriptor = tableInfo.toTableDescriptor();
+        Map<String, String> updateProperties =
+                new HashMap<>(existingTableDescriptor.getProperties());
+        Map<String, String> updateCustomProperties =
+                new HashMap<>(existingTableDescriptor.getCustomProperties());
+        updateProperties.put("table.datalake.enabled", "true");
+        updateCustomProperties.put("table.datalake.enabled", "true");
+
+        TableDescriptor newTableDescriptor =
+                TableDescriptor.builder()
+                        .schema(existingTableDescriptor.getSchema())
+                        .comment(existingTableDescriptor.getComment().orElse("test table"))
+                        .partitionedBy(existingTableDescriptor.getPartitionKeys())
+                        .distributedBy(
+                                existingTableDescriptor
+                                        .getTableDistribution()
+                                        .get()
+                                        .getBucketCount()
+                                        .orElse(3),
+                                existingTableDescriptor.getBucketKeys())
+                        .properties(updateProperties)
+                        .customProperties(updateCustomProperties)
+                        .build();
+        // alter table
+        admin.alterTable(tablePath, newTableDescriptor, false).get();
+
+        TableInfo alteredTableInfo = admin.getTableInfo(tablePath).get();
+        TableDescriptor alteredTableDescriptor = alteredTableInfo.toTableDescriptor();
+        assertThat(alteredTableDescriptor).isEqualTo(newTableDescriptor);
+
+        // throw exception if table not exist
+        assertThatThrownBy(
+                        () ->
+                                admin.alterTable(
+                                                TablePath.of("test_db", "alter_table_not_exist"),
+                                                newTableDescriptor,
+                                                false)
+                                        .get())
+                .cause()
+                .isInstanceOf(TableNotExistException.class);
+
+        // throw exception if database not exist
+        assertThatThrownBy(
+                        () ->
+                                admin.alterTable(
+                                                TablePath.of(
+                                                        "test_db_not_exist",
+                                                        "alter_table_not_exist"),
+                                                newTableDescriptor,
+                                                false)
+                                        .get())
+                .cause()
+                .isInstanceOf(DatabaseNotExistException.class);
+    }
+
     @Test
     void testCreateInvalidDatabaseAndTable() {
         assertThatThrownBy(

fluss-client/src/test/java/org/apache/fluss/client/table/FlussTableITCase.java

@@ -645,6 +645,61 @@ void testPutAndPoll(String kvFormat) throws Exception {
         verifyAppendOrPut(false, "ARROW", kvFormat);
     }
 
+    @Test
+    void testAlterTableBucket() throws Exception {
+        TableDescriptor data1TableDescriptor =
+                TableDescriptor.builder().schema(DATA1_SCHEMA).distributedBy(1).build();
+        createTable(DATA1_TABLE_PATH, data1TableDescriptor, false);
+
+        long tableId;
+        try (Table table = conn.getTable(DATA1_TABLE_PATH)) {
+            tableId = table.getTableInfo().getTableId();
+
+            AppendWriter appendWriter = table.newAppend().createWriter();
+
+            for (int i = 0; i < 10; i++) {
+                appendWriter.append(row(i, "a")).get();
+            }
+
+            try (LogScanner logScanner = createLogScanner(table)) {
+                subscribeFromBeginning(logScanner, table);
+
+                ScanRecords scanRecords = logScanner.poll(Duration.ofSeconds(1));
+                for (ScanRecord record : scanRecords) {
+                    System.out.println(record);
+                }
+            }
+        }
+
+        System.out.println("==");
+
+        TableDescriptor data2TableDescriptor =
+                TableDescriptor.builder().schema(DATA1_SCHEMA).distributedBy(2).build();
+        admin.alterTable(DATA1_TABLE_PATH, data2TableDescriptor, false);
+
+        waitAllReplicasReady(tableId, 2);
+
+        conn = ConnectionFactory.createConnection(clientConf);
+        try (Table table = conn.getTable(DATA1_TABLE_PATH)) {
+            AppendWriter appendWriter = table.newAppend().createWriter();
+
+            for (int i = 0; i < 10; i++) {
+                appendWriter.append(row(i, "a")).get();
+            }
+
+            try (LogScanner logScanner = createLogScanner(table)) {
+                subscribeFromBeginning(logScanner, table);
+
+                ScanRecords scanRecords = logScanner.poll(Duration.ofSeconds(1));
+                for (TableBucket tableBucket : scanRecords.buckets()) {
+                    for (ScanRecord record : scanRecords.records(tableBucket)) {
+                        System.out.println("Bucket: " + tableBucket + " record: " + record);
+                    }
+                }
+            }
+        }
+    }
+
     void verifyAppendOrPut(boolean append, String logFormat, @Nullable String kvFormat)
             throws Exception {
         Schema schema =