You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@iotdb.apache.org by "OneSizeFitsQuorum (via GitHub)" <gi...@apache.org> on 2023/07/31 07:21:37 UTC

[GitHub] [iotdb] OneSizeFitsQuorum commented on a diff in pull request #10648: [IOTDB-6079] Cluster computing resource balance

OneSizeFitsQuorum commented on code in PR #10648:
URL: https://github.com/apache/iotdb/pull/10648#discussion_r1278773349


##########
integration-test/src/test/java/org/apache/iotdb/session/it/IoTDBSessionSchemaTemplateIT.java:
##########
@@ -471,18 +471,22 @@ public void testHybridAutoExtendSchemaTemplate()
 
     Assert.assertTrue(expectedSeries.isEmpty());
 
-    deviceIds = Arrays.asList("root.db.v4.d1", "root.db.v4.d2");
-    timestamps = Arrays.asList(1L, 1L);
     measurements = Arrays.asList("a", "b", "c");
-    allMeasurements = Arrays.asList(measurements, measurements);
-    allTsDataTypes =
-        Arrays.asList(
-            Arrays.asList(TSDataType.FLOAT, TSDataType.FLOAT, TSDataType.TEXT),
-            Arrays.asList(TSDataType.DOUBLE, TSDataType.DOUBLE, TSDataType.INT32));
-    allValues = Arrays.asList(Arrays.asList(1f, 2f, "3"), Arrays.asList(1d, 2d, 3));
 
     try {
-      session.insertRecords(deviceIds, timestamps, allMeasurements, allTsDataTypes, allValues);
+      session.insertRecord(

Review Comment:
   I'm kind of curious why we change here? Their logic seems to be the same



##########
integration-test/src/test/java/org/apache/iotdb/confignode/it/partition/IoTDBPartitionInheritPolicyIT.java:
##########
@@ -95,76 +91,173 @@ public static void tearDown() {
 
   @Test
   public void testDataPartitionInheritPolicy() throws Exception {
+    final long baseStartTime = 1000;
+    Map<TSeriesPartitionSlot, TConsensusGroupId> dataAllotTable1 = new ConcurrentHashMap<>();
 
-    try (SyncConfigNodeIServiceClient client =
-        (SyncConfigNodeIServiceClient) EnvFactory.getEnv().getLeaderConfigNodeConnection()) {
-      TDataPartitionReq dataPartitionReq = new TDataPartitionReq();
-      TDataPartitionTableResp dataPartitionTableResp;
-      Map<String, Map<TSeriesPartitionSlot, TTimeSlotList>> partitionSlotsMap;
-
-      for (int i = 0; i < storageGroupNum; i++) {
-        String storageGroup = sg + i;
-        for (int j = 0; j < testSeriesPartitionSlotNum; j += seriesPartitionBatchSize) {
-          // Test inherit predecessor or successor
-          boolean isAscending = (j / 10) % 2 == 0;
-          int step = isAscending ? timePartitionBatchSize : -timePartitionBatchSize;
-          int k = isAscending ? 0 : testTimePartitionSlotsNum - timePartitionBatchSize;
-          while (0 <= k && k < testTimePartitionSlotsNum) {
-            partitionSlotsMap =
-                ConfigNodeTestUtils.constructPartitionSlotsMap(
-                    storageGroup,
-                    j,
-                    j + seriesPartitionBatchSize,
-                    k,
-                    k + timePartitionBatchSize,
-                    testTimePartitionInterval);
-            // Let ConfigNode create DataPartition
-            dataPartitionReq.setPartitionSlotsMap(partitionSlotsMap);
-            for (int retry = 0; retry < 5; retry++) {
-              // Build new Client since it's unstable
-              try (SyncConfigNodeIServiceClient configNodeClient =
-                  (SyncConfigNodeIServiceClient)
-                      EnvFactory.getEnv().getLeaderConfigNodeConnection()) {
-                dataPartitionTableResp =
-                    configNodeClient.getOrCreateDataPartitionTable(dataPartitionReq);
-                if (dataPartitionTableResp != null
-                    && dataPartitionTableResp.getStatus().getCode()
-                        == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
-                  ConfigNodeTestUtils.checkDataPartitionTable(
-                      storageGroup,
-                      j,
-                      j + seriesPartitionBatchSize,
-                      k,
-                      k + timePartitionBatchSize,
-                      testTimePartitionInterval,
-                      configNodeClient
-                          .getDataPartitionTable(dataPartitionReq)
-                          .getDataPartitionTable());
-                  break;
-                }
-              } catch (Exception e) {
-                // Retry sometimes in order to avoid request timeout
-                LOGGER.error(e.getMessage());
-                TimeUnit.SECONDS.sleep(1);
+    // Test1: divide and inherit DataPartitions from scratch
+    for (long timePartitionSlot = baseStartTime;
+        timePartitionSlot < baseStartTime + testTimePartitionSlotsNum;
+        timePartitionSlot++) {
+      for (int seriesPartitionSlot = 0;
+          seriesPartitionSlot < testSeriesSlotNum;
+          seriesPartitionSlot += seriesPartitionSlotBatchSize) {
+        ConfigNodeTestUtils.getOrCreateDataPartitionWithRetry(
+            database,
+            seriesPartitionSlot,
+            seriesPartitionSlot + seriesPartitionSlotBatchSize,
+            timePartitionSlot,
+            timePartitionSlot + 1,
+            testTimePartitionInterval);
+      }
+    }
+
+    int mu = (int) (testSeriesSlotNum / testDataRegionPerDataNode);
+    TDataPartitionTableResp dataPartitionTableResp =
+        ConfigNodeTestUtils.getDataPartitionWithRetry(
+            database,
+            0,
+            testSeriesSlotNum,
+            baseStartTime,
+            baseStartTime + testTimePartitionSlotsNum,
+            testTimePartitionInterval);
+    Assert.assertNotNull(dataPartitionTableResp);
+
+    // All DataRegionGroups divide all SeriesSlots evenly
+    final int expectedPartitionNum1 = mu * testTimePartitionSlotsNum;
+    Map<TConsensusGroupId, Integer> counter =
+        ConfigNodeTestUtils.countDataPartition(
+            dataPartitionTableResp.getDataPartitionTable().get(database));
+    counter.forEach((groupId, num) -> Assert.assertEquals(expectedPartitionNum1, num.intValue()));
+
+    // Test DataPartition inherit policy
+    dataPartitionTableResp
+        .getDataPartitionTable()
+        .get(database)
+        .forEach(
+            ((seriesPartitionSlot, timePartitionSlotMap) -> {
+              // All Timeslots belonging to the same SeriesSlot are allocated to the same
+              // DataRegionGroup
+              TConsensusGroupId groupId =
+                  timePartitionSlotMap
+                      .get(new TTimePartitionSlot(baseStartTime * testTimePartitionInterval))
+                      .get(0);
+              timePartitionSlotMap.forEach(
+                  (timePartitionSlot, groupIdList) ->
+                      Assert.assertEquals(groupId, groupIdList.get(0)));
+              dataAllotTable1.put(seriesPartitionSlot, groupId);
+            }));
+
+    // Register a new DataNode to extend DataRegionGroups
+    EnvFactory.getEnv().registerNewDataNode(true);
+
+    // Test2: divide and inherit DataPartitions after extension
+    mu = (int) (testSeriesSlotNum / (testDataRegionPerDataNode * 2));
+    dataPartitionTableResp =
+        ConfigNodeTestUtils.getOrCreateDataPartitionWithRetry(
+            database,
+            0,
+            testSeriesSlotNum,
+            baseStartTime + testTimePartitionSlotsNum,
+            baseStartTime + testTimePartitionSlotsNum + timePartitionBatchSize,
+            testTimePartitionInterval);
+    Assert.assertNotNull(dataPartitionTableResp);
+
+    // All DataRegionGroups divide all SeriesSlots evenly
+    counter =
+        ConfigNodeTestUtils.countDataPartition(
+            dataPartitionTableResp.getDataPartitionTable().get(database));
+    final int expectedPartitionNum2 = mu * timePartitionBatchSize;
+    counter.forEach((groupId, num) -> Assert.assertEquals(expectedPartitionNum2, num.intValue()));
+
+    // Test DataPartition inherit policy
+    Map<TSeriesPartitionSlot, TConsensusGroupId> dataAllotTable2 = new ConcurrentHashMap<>();
+    dataPartitionTableResp
+        .getDataPartitionTable()
+        .get(database)
+        .forEach(
+            ((seriesPartitionSlot, timePartitionSlotMap) -> {
+              // All Timeslots belonging to the same SeriesSlot are allocated to the same
+              // DataRegionGroup
+              TConsensusGroupId groupId =
+                  timePartitionSlotMap
+                      .get(
+                          new TTimePartitionSlot(
+                              (baseStartTime + testTimePartitionSlotsNum)
+                                  * testTimePartitionInterval))
+                      .get(0);
+              timePartitionSlotMap.forEach(
+                  (timePartitionSlot, groupIdList) ->
+                      Assert.assertEquals(groupId, groupIdList.get(0)));
+
+              if (dataAllotTable1.containsValue(groupId)) {
+                // The DataRegionGroup has been inherited

Review Comment:
   add size judgement here? Exactly half of the slots should have changed?



##########
iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/structure/BalanceTreeMap.java:
##########
@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.commons.structure;
+
+import org.apache.iotdb.commons.utils.TestOnly;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.TreeMap;
+
+public class BalanceTreeMap<K, V extends Comparable<V>> {
+
+  private final HashMap<K, V> keyValueMap;
+  private final TreeMap<V, Set<K>> valueKeyMap;
+
+  public BalanceTreeMap() {
+    this.keyValueMap = new HashMap<>();
+    this.valueKeyMap = new TreeMap<>();
+  }
+
+  /**
+   * Put or modify a key-value pair.
+   *
+   * @param key Key
+   * @param value Value
+   */
+  public void put(K key, V value) {
+    V oldValue = keyValueMap.getOrDefault(key, null);
+
+    // Update keyValueMap
+    keyValueMap.put(key, value);

Review Comment:
   you can directly use `V oldValue = keyValueMap.put(key, value);`



##########
iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/PartitionBalancer.java:
##########
@@ -16,58 +16,237 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+
 package org.apache.iotdb.confignode.manager.load.balancer;
 
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
 import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot;
+import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
 import org.apache.iotdb.commons.partition.DataPartitionTable;
 import org.apache.iotdb.commons.partition.SchemaPartitionTable;
+import org.apache.iotdb.commons.partition.SeriesPartitionTable;
+import org.apache.iotdb.commons.structure.BalanceTreeMap;
+import org.apache.iotdb.confignode.exception.DatabaseNotExistsException;
 import org.apache.iotdb.confignode.exception.NoAvailableRegionGroupException;
 import org.apache.iotdb.confignode.manager.IManager;
-import org.apache.iotdb.confignode.manager.load.balancer.partition.GreedyPartitionAllocator;
-import org.apache.iotdb.confignode.manager.load.balancer.partition.IPartitionAllocator;
+import org.apache.iotdb.confignode.manager.load.balancer.partition.DataPartitionPolicyTable;
+import org.apache.iotdb.confignode.manager.partition.PartitionManager;
+import org.apache.iotdb.confignode.manager.schema.ClusterSchemaManager;
 import org.apache.iotdb.confignode.rpc.thrift.TTimeSlotList;
+import org.apache.iotdb.tsfile.utils.Pair;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * The SeriesPartitionSlotBalancer provides interfaces to generate optimal Partition allocation and
  * migration plans
  */
 public class PartitionBalancer {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(PartitionBalancer.class);
+
   private final IManager configManager;
 
+  // Map<DatabaseName, DataPartitionPolicyTable>
+  private final Map<String, DataPartitionPolicyTable> dataAllotTableMap;
+
   public PartitionBalancer(IManager configManager) {
     this.configManager = configManager;
+    this.dataAllotTableMap = new ConcurrentHashMap<>();
   }
 
   /**
    * Allocate SchemaPartitions
    *
    * @param unassignedSchemaPartitionSlotsMap SchemaPartitionSlots that should be assigned
-   * @return Map<StorageGroupName, SchemaPartitionTable>, the allocating result
+   * @return Map<DatabaseName, SchemaPartitionTable>, the allocating result
    */
   public Map<String, SchemaPartitionTable> allocateSchemaPartition(
       Map<String, List<TSeriesPartitionSlot>> unassignedSchemaPartitionSlotsMap)
       throws NoAvailableRegionGroupException {
-    return genPartitionAllocator().allocateSchemaPartition(unassignedSchemaPartitionSlotsMap);
+    Map<String, SchemaPartitionTable> result = new ConcurrentHashMap<>();
+
+    for (Map.Entry<String, List<TSeriesPartitionSlot>> slotsMapEntry :
+        unassignedSchemaPartitionSlotsMap.entrySet()) {
+      final String database = slotsMapEntry.getKey();
+      final List<TSeriesPartitionSlot> unassignedPartitionSlots = slotsMapEntry.getValue();
+
+      // Filter available SchemaRegionGroups and
+      // sort them by the number of allocated SchemaPartitions
+      BalanceTreeMap<TConsensusGroupId, Integer> counter = new BalanceTreeMap<>();
+      List<Pair<Long, TConsensusGroupId>> regionSlotsCounter =
+          getPartitionManager()
+              .getSortedRegionGroupSlotsCounter(database, TConsensusGroupType.SchemaRegion);
+      for (Pair<Long, TConsensusGroupId> pair : regionSlotsCounter) {
+        counter.put(pair.getRight(), pair.getLeft().intValue());
+      }
+
+      // Enumerate SeriesPartitionSlot
+      Map<TSeriesPartitionSlot, TConsensusGroupId> schemaPartitionMap = new ConcurrentHashMap<>();

Review Comment:
   same as above



##########
iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/PartitionBalancer.java:
##########
@@ -16,58 +16,237 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+
 package org.apache.iotdb.confignode.manager.load.balancer;
 
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
 import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot;
+import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
 import org.apache.iotdb.commons.partition.DataPartitionTable;
 import org.apache.iotdb.commons.partition.SchemaPartitionTable;
+import org.apache.iotdb.commons.partition.SeriesPartitionTable;
+import org.apache.iotdb.commons.structure.BalanceTreeMap;
+import org.apache.iotdb.confignode.exception.DatabaseNotExistsException;
 import org.apache.iotdb.confignode.exception.NoAvailableRegionGroupException;
 import org.apache.iotdb.confignode.manager.IManager;
-import org.apache.iotdb.confignode.manager.load.balancer.partition.GreedyPartitionAllocator;
-import org.apache.iotdb.confignode.manager.load.balancer.partition.IPartitionAllocator;
+import org.apache.iotdb.confignode.manager.load.balancer.partition.DataPartitionPolicyTable;
+import org.apache.iotdb.confignode.manager.partition.PartitionManager;
+import org.apache.iotdb.confignode.manager.schema.ClusterSchemaManager;
 import org.apache.iotdb.confignode.rpc.thrift.TTimeSlotList;
+import org.apache.iotdb.tsfile.utils.Pair;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * The SeriesPartitionSlotBalancer provides interfaces to generate optimal Partition allocation and
  * migration plans
  */
 public class PartitionBalancer {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(PartitionBalancer.class);
+
   private final IManager configManager;
 
+  // Map<DatabaseName, DataPartitionPolicyTable>
+  private final Map<String, DataPartitionPolicyTable> dataAllotTableMap;
+
   public PartitionBalancer(IManager configManager) {
     this.configManager = configManager;
+    this.dataAllotTableMap = new ConcurrentHashMap<>();
   }
 
   /**
    * Allocate SchemaPartitions
    *
    * @param unassignedSchemaPartitionSlotsMap SchemaPartitionSlots that should be assigned
-   * @return Map<StorageGroupName, SchemaPartitionTable>, the allocating result
+   * @return Map<DatabaseName, SchemaPartitionTable>, the allocating result
    */
   public Map<String, SchemaPartitionTable> allocateSchemaPartition(
       Map<String, List<TSeriesPartitionSlot>> unassignedSchemaPartitionSlotsMap)
       throws NoAvailableRegionGroupException {
-    return genPartitionAllocator().allocateSchemaPartition(unassignedSchemaPartitionSlotsMap);
+    Map<String, SchemaPartitionTable> result = new ConcurrentHashMap<>();

Review Comment:
   Do we really need `ConcurrentHashMap`?



##########
iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/PartitionBalancer.java:
##########
@@ -16,58 +16,237 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+
 package org.apache.iotdb.confignode.manager.load.balancer;
 
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
 import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot;
+import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
 import org.apache.iotdb.commons.partition.DataPartitionTable;
 import org.apache.iotdb.commons.partition.SchemaPartitionTable;
+import org.apache.iotdb.commons.partition.SeriesPartitionTable;
+import org.apache.iotdb.commons.structure.BalanceTreeMap;
+import org.apache.iotdb.confignode.exception.DatabaseNotExistsException;
 import org.apache.iotdb.confignode.exception.NoAvailableRegionGroupException;
 import org.apache.iotdb.confignode.manager.IManager;
-import org.apache.iotdb.confignode.manager.load.balancer.partition.GreedyPartitionAllocator;
-import org.apache.iotdb.confignode.manager.load.balancer.partition.IPartitionAllocator;
+import org.apache.iotdb.confignode.manager.load.balancer.partition.DataPartitionPolicyTable;
+import org.apache.iotdb.confignode.manager.partition.PartitionManager;
+import org.apache.iotdb.confignode.manager.schema.ClusterSchemaManager;
 import org.apache.iotdb.confignode.rpc.thrift.TTimeSlotList;
+import org.apache.iotdb.tsfile.utils.Pair;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * The SeriesPartitionSlotBalancer provides interfaces to generate optimal Partition allocation and
  * migration plans
  */
 public class PartitionBalancer {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(PartitionBalancer.class);
+
   private final IManager configManager;
 
+  // Map<DatabaseName, DataPartitionPolicyTable>
+  private final Map<String, DataPartitionPolicyTable> dataAllotTableMap;
+
   public PartitionBalancer(IManager configManager) {
     this.configManager = configManager;
+    this.dataAllotTableMap = new ConcurrentHashMap<>();
   }
 
   /**
    * Allocate SchemaPartitions
    *
    * @param unassignedSchemaPartitionSlotsMap SchemaPartitionSlots that should be assigned
-   * @return Map<StorageGroupName, SchemaPartitionTable>, the allocating result
+   * @return Map<DatabaseName, SchemaPartitionTable>, the allocating result
    */
   public Map<String, SchemaPartitionTable> allocateSchemaPartition(
       Map<String, List<TSeriesPartitionSlot>> unassignedSchemaPartitionSlotsMap)
       throws NoAvailableRegionGroupException {
-    return genPartitionAllocator().allocateSchemaPartition(unassignedSchemaPartitionSlotsMap);
+    Map<String, SchemaPartitionTable> result = new ConcurrentHashMap<>();
+
+    for (Map.Entry<String, List<TSeriesPartitionSlot>> slotsMapEntry :
+        unassignedSchemaPartitionSlotsMap.entrySet()) {
+      final String database = slotsMapEntry.getKey();
+      final List<TSeriesPartitionSlot> unassignedPartitionSlots = slotsMapEntry.getValue();
+
+      // Filter available SchemaRegionGroups and
+      // sort them by the number of allocated SchemaPartitions
+      BalanceTreeMap<TConsensusGroupId, Integer> counter = new BalanceTreeMap<>();
+      List<Pair<Long, TConsensusGroupId>> regionSlotsCounter =
+          getPartitionManager()
+              .getSortedRegionGroupSlotsCounter(database, TConsensusGroupType.SchemaRegion);
+      for (Pair<Long, TConsensusGroupId> pair : regionSlotsCounter) {
+        counter.put(pair.getRight(), pair.getLeft().intValue());
+      }
+
+      // Enumerate SeriesPartitionSlot
+      Map<TSeriesPartitionSlot, TConsensusGroupId> schemaPartitionMap = new ConcurrentHashMap<>();
+      for (TSeriesPartitionSlot seriesPartitionSlot : unassignedPartitionSlots) {
+        // Greedy allocation: allocate the unassigned SchemaPartition to
+        // the RegionGroup whose allocated SchemaPartitions is the least
+        TConsensusGroupId consensusGroupId = counter.getKeyWithMinValue();
+        schemaPartitionMap.put(seriesPartitionSlot, consensusGroupId);
+        counter.put(consensusGroupId, counter.get(consensusGroupId) + 1);
+      }
+      result.put(database, new SchemaPartitionTable(schemaPartitionMap));
+    }
+
+    return result;
   }
 
   /**
    * Allocate DataPartitions
    *
    * @param unassignedDataPartitionSlotsMap DataPartitionSlots that should be assigned
-   * @return Map<StorageGroupName, DataPartitionTable>, the allocating result
+   * @return Map<DatabaseName, DataPartitionTable>, the allocating result
    */
   public Map<String, DataPartitionTable> allocateDataPartition(
       Map<String, Map<TSeriesPartitionSlot, TTimeSlotList>> unassignedDataPartitionSlotsMap)
       throws NoAvailableRegionGroupException {
-    return genPartitionAllocator().allocateDataPartition(unassignedDataPartitionSlotsMap);
+    Map<String, DataPartitionTable> result = new ConcurrentHashMap<>();
+
+    for (Map.Entry<String, Map<TSeriesPartitionSlot, TTimeSlotList>> slotsMapEntry :
+        unassignedDataPartitionSlotsMap.entrySet()) {
+      final String database = slotsMapEntry.getKey();
+      final Map<TSeriesPartitionSlot, TTimeSlotList> unassignedPartitionSlotsMap =
+          slotsMapEntry.getValue();
+
+      // Filter available DataRegionGroups and
+      // sort them by the number of allocated DataPartitions
+      BalanceTreeMap<TConsensusGroupId, Integer> counter = new BalanceTreeMap<>();
+      List<Pair<Long, TConsensusGroupId>> regionSlotsCounter =
+          getPartitionManager()
+              .getSortedRegionGroupSlotsCounter(database, TConsensusGroupType.DataRegion);
+      for (Pair<Long, TConsensusGroupId> pair : regionSlotsCounter) {
+        counter.put(pair.getRight(), pair.getLeft().intValue());
+      }
+
+      DataPartitionPolicyTable allotTable = dataAllotTableMap.get(database);
+      allotTable.acquireLock();

Review Comment:
   put this in try scope ?



##########
iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/PartitionBalancer.java:
##########
@@ -16,58 +16,237 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+
 package org.apache.iotdb.confignode.manager.load.balancer;
 
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
 import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot;
+import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
 import org.apache.iotdb.commons.partition.DataPartitionTable;
 import org.apache.iotdb.commons.partition.SchemaPartitionTable;
+import org.apache.iotdb.commons.partition.SeriesPartitionTable;
+import org.apache.iotdb.commons.structure.BalanceTreeMap;
+import org.apache.iotdb.confignode.exception.DatabaseNotExistsException;
 import org.apache.iotdb.confignode.exception.NoAvailableRegionGroupException;
 import org.apache.iotdb.confignode.manager.IManager;
-import org.apache.iotdb.confignode.manager.load.balancer.partition.GreedyPartitionAllocator;
-import org.apache.iotdb.confignode.manager.load.balancer.partition.IPartitionAllocator;
+import org.apache.iotdb.confignode.manager.load.balancer.partition.DataPartitionPolicyTable;
+import org.apache.iotdb.confignode.manager.partition.PartitionManager;
+import org.apache.iotdb.confignode.manager.schema.ClusterSchemaManager;
 import org.apache.iotdb.confignode.rpc.thrift.TTimeSlotList;
+import org.apache.iotdb.tsfile.utils.Pair;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * The SeriesPartitionSlotBalancer provides interfaces to generate optimal Partition allocation and
  * migration plans
  */
 public class PartitionBalancer {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(PartitionBalancer.class);
+
   private final IManager configManager;
 
+  // Map<DatabaseName, DataPartitionPolicyTable>
+  private final Map<String, DataPartitionPolicyTable> dataAllotTableMap;
+
   public PartitionBalancer(IManager configManager) {
     this.configManager = configManager;
+    this.dataAllotTableMap = new ConcurrentHashMap<>();
   }
 
   /**
    * Allocate SchemaPartitions
    *
    * @param unassignedSchemaPartitionSlotsMap SchemaPartitionSlots that should be assigned
-   * @return Map<StorageGroupName, SchemaPartitionTable>, the allocating result
+   * @return Map<DatabaseName, SchemaPartitionTable>, the allocating result
    */
   public Map<String, SchemaPartitionTable> allocateSchemaPartition(
       Map<String, List<TSeriesPartitionSlot>> unassignedSchemaPartitionSlotsMap)
       throws NoAvailableRegionGroupException {
-    return genPartitionAllocator().allocateSchemaPartition(unassignedSchemaPartitionSlotsMap);
+    Map<String, SchemaPartitionTable> result = new ConcurrentHashMap<>();
+
+    for (Map.Entry<String, List<TSeriesPartitionSlot>> slotsMapEntry :
+        unassignedSchemaPartitionSlotsMap.entrySet()) {
+      final String database = slotsMapEntry.getKey();
+      final List<TSeriesPartitionSlot> unassignedPartitionSlots = slotsMapEntry.getValue();
+
+      // Filter available SchemaRegionGroups and
+      // sort them by the number of allocated SchemaPartitions
+      BalanceTreeMap<TConsensusGroupId, Integer> counter = new BalanceTreeMap<>();
+      List<Pair<Long, TConsensusGroupId>> regionSlotsCounter =
+          getPartitionManager()
+              .getSortedRegionGroupSlotsCounter(database, TConsensusGroupType.SchemaRegion);
+      for (Pair<Long, TConsensusGroupId> pair : regionSlotsCounter) {
+        counter.put(pair.getRight(), pair.getLeft().intValue());
+      }
+
+      // Enumerate SeriesPartitionSlot
+      Map<TSeriesPartitionSlot, TConsensusGroupId> schemaPartitionMap = new ConcurrentHashMap<>();
+      for (TSeriesPartitionSlot seriesPartitionSlot : unassignedPartitionSlots) {
+        // Greedy allocation: allocate the unassigned SchemaPartition to
+        // the RegionGroup whose allocated SchemaPartitions is the least
+        TConsensusGroupId consensusGroupId = counter.getKeyWithMinValue();
+        schemaPartitionMap.put(seriesPartitionSlot, consensusGroupId);
+        counter.put(consensusGroupId, counter.get(consensusGroupId) + 1);
+      }
+      result.put(database, new SchemaPartitionTable(schemaPartitionMap));
+    }
+
+    return result;
   }
 
   /**
    * Allocate DataPartitions
    *
    * @param unassignedDataPartitionSlotsMap DataPartitionSlots that should be assigned
-   * @return Map<StorageGroupName, DataPartitionTable>, the allocating result
+   * @return Map<DatabaseName, DataPartitionTable>, the allocating result
    */
   public Map<String, DataPartitionTable> allocateDataPartition(
       Map<String, Map<TSeriesPartitionSlot, TTimeSlotList>> unassignedDataPartitionSlotsMap)
       throws NoAvailableRegionGroupException {
-    return genPartitionAllocator().allocateDataPartition(unassignedDataPartitionSlotsMap);
+    Map<String, DataPartitionTable> result = new ConcurrentHashMap<>();

Review Comment:
   same as above



##########
iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/partition/PartitionInfo.java:
##########
@@ -784,6 +822,19 @@ public Set<Integer> getAllSchemaPartition() {
     return schemaPartitionSet;
   }
 
+  /**
+   * Get the last DataAllotTable of the specified Database.
+   *
+   * @param database The specified Database
+   * @return The last DataAllotTable
+   */
+  public Map<TSeriesPartitionSlot, TConsensusGroupId> getLastDataAllotTable(String database) {
+    if (isDatabaseExisted(database)) {
+      return databasePartitionTables.get(database).getLastDataAllotTable();
+    }
+    return new HashMap<>();

Review Comment:
   Collections.emptyMap()



##########
integration-test/src/test/java/org/apache/iotdb/confignode/it/partition/IoTDBPartitionCreationIT.java:
##########
@@ -419,182 +419,6 @@ public void testPartitionAllocation() throws Exception {
           4 + testTimePartitionBatchSize,
           testTimePartitionInterval,
           dataPartitionTableResp.getDataPartitionTable());
-

Review Comment:
   Why we remove these ITs? Do we still have IT in the exception state of region



##########
iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/partition/DataPartitionEntry.java:
##########
@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.commons.partition;
+
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot;
+import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
+
+import java.security.SecureRandom;
+import java.util.Objects;
+
+public class DataPartitionEntry implements Comparable<DataPartitionEntry> {

Review Comment:
   It doesn't look like this class is being used, so why does it exist?



##########
iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/structure/BalanceTreeMap.java:
##########
@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.commons.structure;
+
+import org.apache.iotdb.commons.utils.TestOnly;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.TreeMap;
+
+public class BalanceTreeMap<K, V extends Comparable<V>> {
+
+  private final HashMap<K, V> keyValueMap;
+  private final TreeMap<V, Set<K>> valueKeyMap;
+
+  public BalanceTreeMap() {
+    this.keyValueMap = new HashMap<>();
+    this.valueKeyMap = new TreeMap<>();
+  }
+
+  /**
+   * Put or modify a key-value pair.
+   *
+   * @param key Key
+   * @param value Value
+   */
+  public void put(K key, V value) {
+    V oldValue = keyValueMap.getOrDefault(key, null);
+
+    // Update keyValueMap
+    keyValueMap.put(key, value);
+
+    // Update valueKeyMap
+    if (oldValue != null) {
+      valueKeyMap.get(oldValue).remove(key);
+      if (valueKeyMap.get(oldValue).isEmpty()) {
+        valueKeyMap.remove(oldValue);
+      }
+    }
+    valueKeyMap.computeIfAbsent(value, empty -> new HashSet<>()).add(key);
+  }
+
+  /**
+   * Get key with minimum value.
+   *
+   * @return Key with minimum value
+   */
+  public K getKeyWithMinValue() {
+    return valueKeyMap.firstEntry().getValue().iterator().next();
+  }
+
+  public V get(K key) {
+    return keyValueMap.getOrDefault(key, null);
+  }
+
+  public boolean containsKey(K key) {
+    return keyValueMap.containsKey(key);
+  }
+
+  public int size() {
+    return keyValueMap.size();
+  }
+
+  @TestOnly
+  public void remove(K key) {

Review Comment:
   ```
   @TestOnly
   public void remove(K key) {
       V value = keyValueMap.remove(key);
       if (value != null) {
           Set<K> keysSet = valueKeyMap.get(value);
           if (keysSet != null) {
               keysSet.remove(key);
               if (keysSet.isEmpty()) {
                   valueKeyMap.remove(value);
               }
           }
       }
   }
   ```



##########
integration-test/src/test/java/org/apache/iotdb/confignode/it/partition/IoTDBPartitionInheritPolicyIT.java:
##########
@@ -95,76 +91,173 @@ public static void tearDown() {
 
   @Test
   public void testDataPartitionInheritPolicy() throws Exception {
+    final long baseStartTime = 1000;
+    Map<TSeriesPartitionSlot, TConsensusGroupId> dataAllotTable1 = new ConcurrentHashMap<>();
 
-    try (SyncConfigNodeIServiceClient client =
-        (SyncConfigNodeIServiceClient) EnvFactory.getEnv().getLeaderConfigNodeConnection()) {
-      TDataPartitionReq dataPartitionReq = new TDataPartitionReq();
-      TDataPartitionTableResp dataPartitionTableResp;
-      Map<String, Map<TSeriesPartitionSlot, TTimeSlotList>> partitionSlotsMap;
-
-      for (int i = 0; i < storageGroupNum; i++) {
-        String storageGroup = sg + i;
-        for (int j = 0; j < testSeriesPartitionSlotNum; j += seriesPartitionBatchSize) {
-          // Test inherit predecessor or successor
-          boolean isAscending = (j / 10) % 2 == 0;
-          int step = isAscending ? timePartitionBatchSize : -timePartitionBatchSize;
-          int k = isAscending ? 0 : testTimePartitionSlotsNum - timePartitionBatchSize;
-          while (0 <= k && k < testTimePartitionSlotsNum) {
-            partitionSlotsMap =
-                ConfigNodeTestUtils.constructPartitionSlotsMap(
-                    storageGroup,
-                    j,
-                    j + seriesPartitionBatchSize,
-                    k,
-                    k + timePartitionBatchSize,
-                    testTimePartitionInterval);
-            // Let ConfigNode create DataPartition
-            dataPartitionReq.setPartitionSlotsMap(partitionSlotsMap);
-            for (int retry = 0; retry < 5; retry++) {
-              // Build new Client since it's unstable
-              try (SyncConfigNodeIServiceClient configNodeClient =
-                  (SyncConfigNodeIServiceClient)
-                      EnvFactory.getEnv().getLeaderConfigNodeConnection()) {
-                dataPartitionTableResp =
-                    configNodeClient.getOrCreateDataPartitionTable(dataPartitionReq);
-                if (dataPartitionTableResp != null
-                    && dataPartitionTableResp.getStatus().getCode()
-                        == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
-                  ConfigNodeTestUtils.checkDataPartitionTable(
-                      storageGroup,
-                      j,
-                      j + seriesPartitionBatchSize,
-                      k,
-                      k + timePartitionBatchSize,
-                      testTimePartitionInterval,
-                      configNodeClient
-                          .getDataPartitionTable(dataPartitionReq)
-                          .getDataPartitionTable());
-                  break;
-                }
-              } catch (Exception e) {
-                // Retry sometimes in order to avoid request timeout
-                LOGGER.error(e.getMessage());
-                TimeUnit.SECONDS.sleep(1);
+    // Test1: divide and inherit DataPartitions from scratch
+    for (long timePartitionSlot = baseStartTime;
+        timePartitionSlot < baseStartTime + testTimePartitionSlotsNum;
+        timePartitionSlot++) {
+      for (int seriesPartitionSlot = 0;
+          seriesPartitionSlot < testSeriesSlotNum;
+          seriesPartitionSlot += seriesPartitionSlotBatchSize) {
+        ConfigNodeTestUtils.getOrCreateDataPartitionWithRetry(
+            database,
+            seriesPartitionSlot,
+            seriesPartitionSlot + seriesPartitionSlotBatchSize,
+            timePartitionSlot,
+            timePartitionSlot + 1,
+            testTimePartitionInterval);
+      }
+    }
+
+    int mu = (int) (testSeriesSlotNum / testDataRegionPerDataNode);
+    TDataPartitionTableResp dataPartitionTableResp =
+        ConfigNodeTestUtils.getDataPartitionWithRetry(
+            database,
+            0,
+            testSeriesSlotNum,
+            baseStartTime,
+            baseStartTime + testTimePartitionSlotsNum,
+            testTimePartitionInterval);
+    Assert.assertNotNull(dataPartitionTableResp);
+
+    // All DataRegionGroups divide all SeriesSlots evenly
+    final int expectedPartitionNum1 = mu * testTimePartitionSlotsNum;
+    Map<TConsensusGroupId, Integer> counter =
+        ConfigNodeTestUtils.countDataPartition(
+            dataPartitionTableResp.getDataPartitionTable().get(database));
+    counter.forEach((groupId, num) -> Assert.assertEquals(expectedPartitionNum1, num.intValue()));
+
+    // Test DataPartition inherit policy
+    dataPartitionTableResp
+        .getDataPartitionTable()
+        .get(database)
+        .forEach(
+            ((seriesPartitionSlot, timePartitionSlotMap) -> {
+              // All Timeslots belonging to the same SeriesSlot are allocated to the same
+              // DataRegionGroup
+              TConsensusGroupId groupId =
+                  timePartitionSlotMap
+                      .get(new TTimePartitionSlot(baseStartTime * testTimePartitionInterval))
+                      .get(0);
+              timePartitionSlotMap.forEach(
+                  (timePartitionSlot, groupIdList) ->
+                      Assert.assertEquals(groupId, groupIdList.get(0)));
+              dataAllotTable1.put(seriesPartitionSlot, groupId);
+            }));
+
+    // Register a new DataNode to extend DataRegionGroups
+    EnvFactory.getEnv().registerNewDataNode(true);
+
+    // Test2: divide and inherit DataPartitions after extension
+    mu = (int) (testSeriesSlotNum / (testDataRegionPerDataNode * 2));
+    dataPartitionTableResp =
+        ConfigNodeTestUtils.getOrCreateDataPartitionWithRetry(
+            database,
+            0,
+            testSeriesSlotNum,
+            baseStartTime + testTimePartitionSlotsNum,
+            baseStartTime + testTimePartitionSlotsNum + timePartitionBatchSize,
+            testTimePartitionInterval);
+    Assert.assertNotNull(dataPartitionTableResp);
+
+    // All DataRegionGroups divide all SeriesSlots evenly
+    counter =
+        ConfigNodeTestUtils.countDataPartition(
+            dataPartitionTableResp.getDataPartitionTable().get(database));
+    final int expectedPartitionNum2 = mu * timePartitionBatchSize;
+    counter.forEach((groupId, num) -> Assert.assertEquals(expectedPartitionNum2, num.intValue()));
+
+    // Test DataPartition inherit policy
+    Map<TSeriesPartitionSlot, TConsensusGroupId> dataAllotTable2 = new ConcurrentHashMap<>();
+    dataPartitionTableResp
+        .getDataPartitionTable()
+        .get(database)
+        .forEach(
+            ((seriesPartitionSlot, timePartitionSlotMap) -> {
+              // All Timeslots belonging to the same SeriesSlot are allocated to the same
+              // DataRegionGroup
+              TConsensusGroupId groupId =
+                  timePartitionSlotMap
+                      .get(
+                          new TTimePartitionSlot(
+                              (baseStartTime + testTimePartitionSlotsNum)
+                                  * testTimePartitionInterval))
+                      .get(0);
+              timePartitionSlotMap.forEach(
+                  (timePartitionSlot, groupIdList) ->
+                      Assert.assertEquals(groupId, groupIdList.get(0)));
+
+              if (dataAllotTable1.containsValue(groupId)) {
+                // The DataRegionGroup has been inherited
+                Assert.assertTrue(dataAllotTable1.containsKey(seriesPartitionSlot));
+                Assert.assertEquals(dataAllotTable1.get(seriesPartitionSlot), groupId);
               }
-            }
-            k += step;
-          }
-        }
+              dataAllotTable2.put(seriesPartitionSlot, groupId);
+            }));
+
+    // Test3: historical DataPartitions will inherit successor
+    Random random = new Random();
+    Set<Integer> allocatedSlots = new HashSet<>();
+    for (int i = 0; i < 10; i++) {
+      int slot = random.nextInt(testSeriesSlotNum);
+      while (allocatedSlots.contains(slot)) {
+        slot = random.nextInt(testSeriesSlotNum);
+      }
+      allocatedSlots.add(slot);
+      dataPartitionTableResp =
+          ConfigNodeTestUtils.getOrCreateDataPartitionWithRetry(
+              database,
+              slot,
+              slot + 1,
+              baseStartTime - 1,
+              baseStartTime,
+              testTimePartitionInterval);
+      Assert.assertNotNull(dataPartitionTableResp);
+
+      TSeriesPartitionSlot seriesPartitionSlot = new TSeriesPartitionSlot(slot);
+      TTimePartitionSlot timePartitionSlot =
+          new TTimePartitionSlot((baseStartTime - 1) * testTimePartitionInterval);
+      Assert.assertEquals(
+          dataAllotTable1.get(seriesPartitionSlot),
+          dataPartitionTableResp
+              .getDataPartitionTable()
+              .get(database)
+              .get(seriesPartitionSlot)
+              .get(timePartitionSlot)
+              .get(0));
+    }
+
+    // Test4: future DataPartitions will inherit predecessor
+    allocatedSlots.clear();
+    for (int i = 0; i < 10; i++) {
+      int slot = random.nextInt(testSeriesSlotNum);
+      while (allocatedSlots.contains(slot)) {
+        slot = random.nextInt(testSeriesSlotNum);
       }
+      allocatedSlots.add(slot);
+      dataPartitionTableResp =
+          ConfigNodeTestUtils.getOrCreateDataPartitionWithRetry(
+              database,
+              slot,
+              slot + 1,
+              baseStartTime + 999,
+              baseStartTime + 1000,
+              testTimePartitionInterval);
+      Assert.assertNotNull(dataPartitionTableResp);
 
-      // Test DataPartition inherit policy
-      TShowRegionResp showRegionResp = client.showRegion(new TShowRegionReq());
-      showRegionResp
-          .getRegionInfoList()
-          .forEach(
-              regionInfo -> {
-                // All Timeslots belonging to the same SeriesSlot are allocated to the same
-                // DataRegionGroup
-                Assert.assertEquals(
-                    regionInfo.getSeriesSlots() * testTimePartitionSlotsNum,
-                    regionInfo.getTimeSlots());
-              });
+      TSeriesPartitionSlot seriesPartitionSlot = new TSeriesPartitionSlot(slot);
+      TTimePartitionSlot timePartitionSlot =
+          new TTimePartitionSlot((baseStartTime + 999) * testTimePartitionInterval);

Review Comment:
   same



##########
iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/structure/BalanceTreeMap.java:
##########
@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.commons.structure;
+
+import org.apache.iotdb.commons.utils.TestOnly;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.TreeMap;
+
+public class BalanceTreeMap<K, V extends Comparable<V>> {
+
+  private final HashMap<K, V> keyValueMap;
+  private final TreeMap<V, Set<K>> valueKeyMap;
+
+  public BalanceTreeMap() {
+    this.keyValueMap = new HashMap<>();
+    this.valueKeyMap = new TreeMap<>();
+  }
+
+  /**
+   * Put or modify a key-value pair.
+   *
+   * @param key Key
+   * @param value Value
+   */
+  public void put(K key, V value) {
+    V oldValue = keyValueMap.getOrDefault(key, null);
+
+    // Update keyValueMap
+    keyValueMap.put(key, value);
+
+    // Update valueKeyMap
+    if (oldValue != null) {

Review Comment:
   ```
     Set<KeyType> keysSet = valueKeyMap.get(oldValue);
     keysSet.remove(key);
     if (keysSet.isEmpty()) {
       valueKeyMap.remove(oldValue);
     }
   ```



##########
integration-test/src/test/java/org/apache/iotdb/confignode/it/partition/IoTDBPartitionInheritPolicyIT.java:
##########
@@ -95,76 +91,173 @@ public static void tearDown() {
 
   @Test
   public void testDataPartitionInheritPolicy() throws Exception {
+    final long baseStartTime = 1000;
+    Map<TSeriesPartitionSlot, TConsensusGroupId> dataAllotTable1 = new ConcurrentHashMap<>();
 
-    try (SyncConfigNodeIServiceClient client =
-        (SyncConfigNodeIServiceClient) EnvFactory.getEnv().getLeaderConfigNodeConnection()) {
-      TDataPartitionReq dataPartitionReq = new TDataPartitionReq();
-      TDataPartitionTableResp dataPartitionTableResp;
-      Map<String, Map<TSeriesPartitionSlot, TTimeSlotList>> partitionSlotsMap;
-
-      for (int i = 0; i < storageGroupNum; i++) {
-        String storageGroup = sg + i;
-        for (int j = 0; j < testSeriesPartitionSlotNum; j += seriesPartitionBatchSize) {
-          // Test inherit predecessor or successor
-          boolean isAscending = (j / 10) % 2 == 0;
-          int step = isAscending ? timePartitionBatchSize : -timePartitionBatchSize;
-          int k = isAscending ? 0 : testTimePartitionSlotsNum - timePartitionBatchSize;
-          while (0 <= k && k < testTimePartitionSlotsNum) {
-            partitionSlotsMap =
-                ConfigNodeTestUtils.constructPartitionSlotsMap(
-                    storageGroup,
-                    j,
-                    j + seriesPartitionBatchSize,
-                    k,
-                    k + timePartitionBatchSize,
-                    testTimePartitionInterval);
-            // Let ConfigNode create DataPartition
-            dataPartitionReq.setPartitionSlotsMap(partitionSlotsMap);
-            for (int retry = 0; retry < 5; retry++) {
-              // Build new Client since it's unstable
-              try (SyncConfigNodeIServiceClient configNodeClient =
-                  (SyncConfigNodeIServiceClient)
-                      EnvFactory.getEnv().getLeaderConfigNodeConnection()) {
-                dataPartitionTableResp =
-                    configNodeClient.getOrCreateDataPartitionTable(dataPartitionReq);
-                if (dataPartitionTableResp != null
-                    && dataPartitionTableResp.getStatus().getCode()
-                        == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
-                  ConfigNodeTestUtils.checkDataPartitionTable(
-                      storageGroup,
-                      j,
-                      j + seriesPartitionBatchSize,
-                      k,
-                      k + timePartitionBatchSize,
-                      testTimePartitionInterval,
-                      configNodeClient
-                          .getDataPartitionTable(dataPartitionReq)
-                          .getDataPartitionTable());
-                  break;
-                }
-              } catch (Exception e) {
-                // Retry sometimes in order to avoid request timeout
-                LOGGER.error(e.getMessage());
-                TimeUnit.SECONDS.sleep(1);
+    // Test1: divide and inherit DataPartitions from scratch
+    for (long timePartitionSlot = baseStartTime;
+        timePartitionSlot < baseStartTime + testTimePartitionSlotsNum;
+        timePartitionSlot++) {
+      for (int seriesPartitionSlot = 0;
+          seriesPartitionSlot < testSeriesSlotNum;
+          seriesPartitionSlot += seriesPartitionSlotBatchSize) {
+        ConfigNodeTestUtils.getOrCreateDataPartitionWithRetry(
+            database,
+            seriesPartitionSlot,
+            seriesPartitionSlot + seriesPartitionSlotBatchSize,
+            timePartitionSlot,
+            timePartitionSlot + 1,
+            testTimePartitionInterval);
+      }
+    }
+
+    int mu = (int) (testSeriesSlotNum / testDataRegionPerDataNode);
+    TDataPartitionTableResp dataPartitionTableResp =
+        ConfigNodeTestUtils.getDataPartitionWithRetry(
+            database,
+            0,
+            testSeriesSlotNum,
+            baseStartTime,
+            baseStartTime + testTimePartitionSlotsNum,
+            testTimePartitionInterval);
+    Assert.assertNotNull(dataPartitionTableResp);
+
+    // All DataRegionGroups divide all SeriesSlots evenly
+    final int expectedPartitionNum1 = mu * testTimePartitionSlotsNum;
+    Map<TConsensusGroupId, Integer> counter =
+        ConfigNodeTestUtils.countDataPartition(
+            dataPartitionTableResp.getDataPartitionTable().get(database));
+    counter.forEach((groupId, num) -> Assert.assertEquals(expectedPartitionNum1, num.intValue()));
+
+    // Test DataPartition inherit policy
+    dataPartitionTableResp
+        .getDataPartitionTable()
+        .get(database)
+        .forEach(
+            ((seriesPartitionSlot, timePartitionSlotMap) -> {
+              // All Timeslots belonging to the same SeriesSlot are allocated to the same
+              // DataRegionGroup
+              TConsensusGroupId groupId =
+                  timePartitionSlotMap
+                      .get(new TTimePartitionSlot(baseStartTime * testTimePartitionInterval))
+                      .get(0);
+              timePartitionSlotMap.forEach(
+                  (timePartitionSlot, groupIdList) ->
+                      Assert.assertEquals(groupId, groupIdList.get(0)));
+              dataAllotTable1.put(seriesPartitionSlot, groupId);
+            }));
+
+    // Register a new DataNode to extend DataRegionGroups
+    EnvFactory.getEnv().registerNewDataNode(true);
+
+    // Test2: divide and inherit DataPartitions after extension
+    mu = (int) (testSeriesSlotNum / (testDataRegionPerDataNode * 2));
+    dataPartitionTableResp =
+        ConfigNodeTestUtils.getOrCreateDataPartitionWithRetry(
+            database,
+            0,
+            testSeriesSlotNum,
+            baseStartTime + testTimePartitionSlotsNum,
+            baseStartTime + testTimePartitionSlotsNum + timePartitionBatchSize,
+            testTimePartitionInterval);
+    Assert.assertNotNull(dataPartitionTableResp);
+
+    // All DataRegionGroups divide all SeriesSlots evenly
+    counter =
+        ConfigNodeTestUtils.countDataPartition(
+            dataPartitionTableResp.getDataPartitionTable().get(database));
+    final int expectedPartitionNum2 = mu * timePartitionBatchSize;
+    counter.forEach((groupId, num) -> Assert.assertEquals(expectedPartitionNum2, num.intValue()));
+
+    // Test DataPartition inherit policy
+    Map<TSeriesPartitionSlot, TConsensusGroupId> dataAllotTable2 = new ConcurrentHashMap<>();
+    dataPartitionTableResp
+        .getDataPartitionTable()
+        .get(database)
+        .forEach(
+            ((seriesPartitionSlot, timePartitionSlotMap) -> {
+              // All Timeslots belonging to the same SeriesSlot are allocated to the same
+              // DataRegionGroup
+              TConsensusGroupId groupId =
+                  timePartitionSlotMap
+                      .get(
+                          new TTimePartitionSlot(
+                              (baseStartTime + testTimePartitionSlotsNum)
+                                  * testTimePartitionInterval))
+                      .get(0);
+              timePartitionSlotMap.forEach(
+                  (timePartitionSlot, groupIdList) ->
+                      Assert.assertEquals(groupId, groupIdList.get(0)));
+
+              if (dataAllotTable1.containsValue(groupId)) {
+                // The DataRegionGroup has been inherited
+                Assert.assertTrue(dataAllotTable1.containsKey(seriesPartitionSlot));
+                Assert.assertEquals(dataAllotTable1.get(seriesPartitionSlot), groupId);
               }
-            }
-            k += step;
-          }
-        }
+              dataAllotTable2.put(seriesPartitionSlot, groupId);
+            }));
+
+    // Test3: historical DataPartitions will inherit successor
+    Random random = new Random();
+    Set<Integer> allocatedSlots = new HashSet<>();
+    for (int i = 0; i < 10; i++) {
+      int slot = random.nextInt(testSeriesSlotNum);
+      while (allocatedSlots.contains(slot)) {
+        slot = random.nextInt(testSeriesSlotNum);
+      }
+      allocatedSlots.add(slot);
+      dataPartitionTableResp =
+          ConfigNodeTestUtils.getOrCreateDataPartitionWithRetry(
+              database,
+              slot,
+              slot + 1,
+              baseStartTime - 1,
+              baseStartTime,
+              testTimePartitionInterval);
+      Assert.assertNotNull(dataPartitionTableResp);
+
+      TSeriesPartitionSlot seriesPartitionSlot = new TSeriesPartitionSlot(slot);
+      TTimePartitionSlot timePartitionSlot =
+          new TTimePartitionSlot((baseStartTime - 1) * testTimePartitionInterval);
+      Assert.assertEquals(
+          dataAllotTable1.get(seriesPartitionSlot),
+          dataPartitionTableResp
+              .getDataPartitionTable()
+              .get(database)
+              .get(seriesPartitionSlot)
+              .get(timePartitionSlot)
+              .get(0));
+    }
+
+    // Test4: future DataPartitions will inherit predecessor
+    allocatedSlots.clear();
+    for (int i = 0; i < 10; i++) {
+      int slot = random.nextInt(testSeriesSlotNum);
+      while (allocatedSlots.contains(slot)) {
+        slot = random.nextInt(testSeriesSlotNum);
       }
+      allocatedSlots.add(slot);
+      dataPartitionTableResp =
+          ConfigNodeTestUtils.getOrCreateDataPartitionWithRetry(
+              database,
+              slot,
+              slot + 1,
+              baseStartTime + 999,

Review Comment:
   do not use magic number, you can use `testTimePartitionSlotsNum * 10 - 1` and `testTimePartitionSlotsNum * 10`



##########
iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/structure/BalanceTreeMap.java:
##########
@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.commons.structure;
+
+import org.apache.iotdb.commons.utils.TestOnly;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.TreeMap;
+
+public class BalanceTreeMap<K, V extends Comparable<V>> {

Review Comment:
   Add some comments for this structure



##########
iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/structure/BalanceTreeMap.java:
##########
@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.commons.structure;
+
+import org.apache.iotdb.commons.utils.TestOnly;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.TreeMap;
+
+public class BalanceTreeMap<K, V extends Comparable<V>> {
+
+  private final HashMap<K, V> keyValueMap;
+  private final TreeMap<V, Set<K>> valueKeyMap;

Review Comment:
   ` valueKeysMap`



##########
iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/PartitionBalancer.java:
##########
@@ -16,58 +16,237 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+
 package org.apache.iotdb.confignode.manager.load.balancer;
 
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
 import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot;
+import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
 import org.apache.iotdb.commons.partition.DataPartitionTable;
 import org.apache.iotdb.commons.partition.SchemaPartitionTable;
+import org.apache.iotdb.commons.partition.SeriesPartitionTable;
+import org.apache.iotdb.commons.structure.BalanceTreeMap;
+import org.apache.iotdb.confignode.exception.DatabaseNotExistsException;
 import org.apache.iotdb.confignode.exception.NoAvailableRegionGroupException;
 import org.apache.iotdb.confignode.manager.IManager;
-import org.apache.iotdb.confignode.manager.load.balancer.partition.GreedyPartitionAllocator;
-import org.apache.iotdb.confignode.manager.load.balancer.partition.IPartitionAllocator;
+import org.apache.iotdb.confignode.manager.load.balancer.partition.DataPartitionPolicyTable;
+import org.apache.iotdb.confignode.manager.partition.PartitionManager;
+import org.apache.iotdb.confignode.manager.schema.ClusterSchemaManager;
 import org.apache.iotdb.confignode.rpc.thrift.TTimeSlotList;
+import org.apache.iotdb.tsfile.utils.Pair;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * The SeriesPartitionSlotBalancer provides interfaces to generate optimal Partition allocation and
  * migration plans
  */
 public class PartitionBalancer {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(PartitionBalancer.class);
+
   private final IManager configManager;
 
+  // Map<DatabaseName, DataPartitionPolicyTable>
+  private final Map<String, DataPartitionPolicyTable> dataAllotTableMap;
+
   public PartitionBalancer(IManager configManager) {
     this.configManager = configManager;
+    this.dataAllotTableMap = new ConcurrentHashMap<>();
   }
 
   /**
    * Allocate SchemaPartitions
    *
    * @param unassignedSchemaPartitionSlotsMap SchemaPartitionSlots that should be assigned
-   * @return Map<StorageGroupName, SchemaPartitionTable>, the allocating result
+   * @return Map<DatabaseName, SchemaPartitionTable>, the allocating result
    */
   public Map<String, SchemaPartitionTable> allocateSchemaPartition(
       Map<String, List<TSeriesPartitionSlot>> unassignedSchemaPartitionSlotsMap)
       throws NoAvailableRegionGroupException {
-    return genPartitionAllocator().allocateSchemaPartition(unassignedSchemaPartitionSlotsMap);
+    Map<String, SchemaPartitionTable> result = new ConcurrentHashMap<>();
+
+    for (Map.Entry<String, List<TSeriesPartitionSlot>> slotsMapEntry :
+        unassignedSchemaPartitionSlotsMap.entrySet()) {
+      final String database = slotsMapEntry.getKey();
+      final List<TSeriesPartitionSlot> unassignedPartitionSlots = slotsMapEntry.getValue();
+
+      // Filter available SchemaRegionGroups and
+      // sort them by the number of allocated SchemaPartitions
+      BalanceTreeMap<TConsensusGroupId, Integer> counter = new BalanceTreeMap<>();
+      List<Pair<Long, TConsensusGroupId>> regionSlotsCounter =
+          getPartitionManager()
+              .getSortedRegionGroupSlotsCounter(database, TConsensusGroupType.SchemaRegion);
+      for (Pair<Long, TConsensusGroupId> pair : regionSlotsCounter) {
+        counter.put(pair.getRight(), pair.getLeft().intValue());
+      }
+
+      // Enumerate SeriesPartitionSlot
+      Map<TSeriesPartitionSlot, TConsensusGroupId> schemaPartitionMap = new ConcurrentHashMap<>();
+      for (TSeriesPartitionSlot seriesPartitionSlot : unassignedPartitionSlots) {
+        // Greedy allocation: allocate the unassigned SchemaPartition to
+        // the RegionGroup whose allocated SchemaPartitions is the least
+        TConsensusGroupId consensusGroupId = counter.getKeyWithMinValue();
+        schemaPartitionMap.put(seriesPartitionSlot, consensusGroupId);
+        counter.put(consensusGroupId, counter.get(consensusGroupId) + 1);
+      }
+      result.put(database, new SchemaPartitionTable(schemaPartitionMap));
+    }
+
+    return result;
   }
 
   /**
    * Allocate DataPartitions
    *
    * @param unassignedDataPartitionSlotsMap DataPartitionSlots that should be assigned
-   * @return Map<StorageGroupName, DataPartitionTable>, the allocating result
+   * @return Map<DatabaseName, DataPartitionTable>, the allocating result
    */
   public Map<String, DataPartitionTable> allocateDataPartition(
       Map<String, Map<TSeriesPartitionSlot, TTimeSlotList>> unassignedDataPartitionSlotsMap)
       throws NoAvailableRegionGroupException {
-    return genPartitionAllocator().allocateDataPartition(unassignedDataPartitionSlotsMap);
+    Map<String, DataPartitionTable> result = new ConcurrentHashMap<>();
+
+    for (Map.Entry<String, Map<TSeriesPartitionSlot, TTimeSlotList>> slotsMapEntry :
+        unassignedDataPartitionSlotsMap.entrySet()) {
+      final String database = slotsMapEntry.getKey();
+      final Map<TSeriesPartitionSlot, TTimeSlotList> unassignedPartitionSlotsMap =
+          slotsMapEntry.getValue();
+
+      // Filter available DataRegionGroups and
+      // sort them by the number of allocated DataPartitions
+      BalanceTreeMap<TConsensusGroupId, Integer> counter = new BalanceTreeMap<>();
+      List<Pair<Long, TConsensusGroupId>> regionSlotsCounter =
+          getPartitionManager()
+              .getSortedRegionGroupSlotsCounter(database, TConsensusGroupType.DataRegion);
+      for (Pair<Long, TConsensusGroupId> pair : regionSlotsCounter) {
+        counter.put(pair.getRight(), pair.getLeft().intValue());
+      }
+
+      DataPartitionPolicyTable allotTable = dataAllotTableMap.get(database);
+      allotTable.acquireLock();
+      DataPartitionTable dataPartitionTable = new DataPartitionTable();
+
+      try {
+        // Enumerate SeriesPartitionSlot
+        for (Map.Entry<TSeriesPartitionSlot, TTimeSlotList> seriesPartitionEntry :
+            unassignedPartitionSlotsMap.entrySet()) {
+          SeriesPartitionTable seriesPartitionTable = new SeriesPartitionTable();
+
+          // Enumerate TimePartitionSlot in ascending order
+          TSeriesPartitionSlot seriesPartitionSlot = seriesPartitionEntry.getKey();
+          List<TTimePartitionSlot> timePartitionSlots =
+              seriesPartitionEntry.getValue().getTimePartitionSlots();
+          timePartitionSlots.sort(Comparator.comparingLong(TTimePartitionSlot::getStartTime));
+
+          for (TTimePartitionSlot timePartitionSlot : timePartitionSlots) {
+
+            // 1. The historical DataPartition will try to inherit successor DataPartition first
+            TConsensusGroupId successor =
+                getPartitionManager()
+                    .getSuccessorDataPartition(database, seriesPartitionSlot, timePartitionSlot);
+            if (successor != null && counter.containsKey(successor)) {

Review Comment:
   Why we need `counter.containsKey(successor)`? Are only running RegionGroups allowed to exist inside them, and how is this guaranteed?



##########
iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/schema/DeleteDatabaseProcedure.java:
##########
@@ -206,6 +206,7 @@ protected Flow executeFromState(ConfigNodeProcedureEnv env, DeleteStorageGroupSt
             LOG.info(
                 "[DeleteDatabaseProcedure] Database: {} is deleted successfully",
                 deleteDatabaseSchema.getName());
+            env.getConfigManager().getLoadManager().clearPartitionBalancer();

Review Comment:
   Whether only the corresponding database table should be clear?



##########
iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/partition/DataPartitionPolicyTable.java:
##########
@@ -0,0 +1,144 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.confignode.manager.load.balancer.partition;
+
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot;
+import org.apache.iotdb.commons.structure.BalanceTreeMap;
+import org.apache.iotdb.confignode.conf.ConfigNodeConfig;
+import org.apache.iotdb.confignode.conf.ConfigNodeDescriptor;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.locks.ReentrantLock;
+
+public class DataPartitionPolicyTable {
+
+  private static final ConfigNodeConfig CONF = ConfigNodeDescriptor.getInstance().getConf();
+  private static final int SERIES_SLOT_NUM = CONF.getSeriesSlotNum();
+
+  private final ReentrantLock dataAllotTableLock;
+
+  // Map<SeriesPartitionSlot, RegionGroupId>
+  // The optimal allocation of SeriesSlots to RegionGroups
+  private final Map<TSeriesPartitionSlot, TConsensusGroupId> dataAllotMap;
+
+  // Map<RegionGroupId, SeriesPartitionSlot Count>
+  // The number of SeriesSlots allocated to each RegionGroup in dataAllotMap
+  private final BalanceTreeMap<TConsensusGroupId, Integer> seriesPartitionSlotCounter;
+
+  public DataPartitionPolicyTable() {
+    this.dataAllotTableLock = new ReentrantLock();
+    this.dataAllotMap = new HashMap<>();
+    this.seriesPartitionSlotCounter = new BalanceTreeMap<>();
+  }
+
+  /**
+   * Get or activate the specified SeriesPartitionSlot in dataAllotMap.
+   *
+   * @param seriesPartitionSlot The specified SeriesPartitionSlot
+   * @return The RegionGroupId of the specified SeriesPartitionSlot, activate when its empty yet
+   */
+  public TConsensusGroupId getRegionGroupIdOrActivateIfNecessary(
+      TSeriesPartitionSlot seriesPartitionSlot) {
+    if (dataAllotMap.containsKey(seriesPartitionSlot)) {
+      return dataAllotMap.get(seriesPartitionSlot);
+    }
+
+    TConsensusGroupId regionGroupId = seriesPartitionSlotCounter.getKeyWithMinValue();
+    dataAllotMap.put(seriesPartitionSlot, regionGroupId);
+    seriesPartitionSlotCounter.put(
+        regionGroupId, seriesPartitionSlotCounter.get(regionGroupId) + 1);
+    return regionGroupId;
+  }
+
+  /**
+   * Re-balance the allocation of SeriesSlots to RegionGroups.
+   *
+   * @param dataRegionGroups All DataRegionGroups currently in the Database
+   */
+  public void reBalanceDataPartitionPolicy(List<TConsensusGroupId> dataRegionGroups) {
+    dataAllotTableLock.lock();
+    try {
+      dataRegionGroups.forEach(
+          dataRegionGroup -> {
+            if (!seriesPartitionSlotCounter.containsKey(dataRegionGroup)) {
+              seriesPartitionSlotCounter.put(dataRegionGroup, 0);
+            }
+          });
+
+      // Enumerate all SeriesSlots randomly
+      List<TSeriesPartitionSlot> seriesPartitionSlots = new ArrayList<>();
+      for (int i = 0; i < SERIES_SLOT_NUM; i++) {
+        seriesPartitionSlots.add(new TSeriesPartitionSlot(i));
+      }
+      Collections.shuffle(seriesPartitionSlots);
+
+      int mu = SERIES_SLOT_NUM / dataRegionGroups.size();
+      for (TSeriesPartitionSlot seriesPartitionSlot : seriesPartitionSlots) {
+        if (!dataAllotMap.containsKey(seriesPartitionSlot)) {
+          continue;
+        }
+
+        TConsensusGroupId regionGroupId = dataAllotMap.get(seriesPartitionSlot);
+        int seriesPartitionSlotCount = seriesPartitionSlotCounter.get(regionGroupId);
+        if (seriesPartitionSlotCount > mu) {
+          // Remove from dataAllotMap if the number of SeriesSlots is greater than mu
+          dataAllotMap.remove(seriesPartitionSlot);
+          seriesPartitionSlotCounter.put(regionGroupId, seriesPartitionSlotCount - 1);
+        }
+      }
+    } finally {
+      dataAllotTableLock.unlock();
+    }
+  }
+
+  /** Only use this interface when init PartitionBalancer. */
+  public void setDataAllotMap(Map<TSeriesPartitionSlot, TConsensusGroupId> dataAllotMap) {
+    dataAllotTableLock.lock();

Review Comment:
   move to the try scope



##########
iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/statemachine/CreateRegionGroupsProcedure.java:
##########
@@ -259,6 +275,7 @@ public void deserialize(ByteBuffer byteBuffer) {
             ThriftCommonsSerDeUtils.deserializeTRegionReplicaSet(byteBuffer);
         failedRegionReplicaSets.put(groupId, replica);
       }
+      persistPlan.deserializeForProcedure(byteBuffer);

Review Comment:
   For compatibility reasons, it is important to determine if there are still redundant bytes that can be read, and if so, deserialize them



##########
iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/PartitionBalancer.java:
##########
@@ -16,58 +16,237 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+
 package org.apache.iotdb.confignode.manager.load.balancer;
 
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
 import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot;
+import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
 import org.apache.iotdb.commons.partition.DataPartitionTable;
 import org.apache.iotdb.commons.partition.SchemaPartitionTable;
+import org.apache.iotdb.commons.partition.SeriesPartitionTable;
+import org.apache.iotdb.commons.structure.BalanceTreeMap;
+import org.apache.iotdb.confignode.exception.DatabaseNotExistsException;
 import org.apache.iotdb.confignode.exception.NoAvailableRegionGroupException;
 import org.apache.iotdb.confignode.manager.IManager;
-import org.apache.iotdb.confignode.manager.load.balancer.partition.GreedyPartitionAllocator;
-import org.apache.iotdb.confignode.manager.load.balancer.partition.IPartitionAllocator;
+import org.apache.iotdb.confignode.manager.load.balancer.partition.DataPartitionPolicyTable;
+import org.apache.iotdb.confignode.manager.partition.PartitionManager;
+import org.apache.iotdb.confignode.manager.schema.ClusterSchemaManager;
 import org.apache.iotdb.confignode.rpc.thrift.TTimeSlotList;
+import org.apache.iotdb.tsfile.utils.Pair;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * The SeriesPartitionSlotBalancer provides interfaces to generate optimal Partition allocation and
  * migration plans
  */
 public class PartitionBalancer {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(PartitionBalancer.class);
+
   private final IManager configManager;
 
+  // Map<DatabaseName, DataPartitionPolicyTable>
+  private final Map<String, DataPartitionPolicyTable> dataAllotTableMap;
+
   public PartitionBalancer(IManager configManager) {
     this.configManager = configManager;
+    this.dataAllotTableMap = new ConcurrentHashMap<>();
   }
 
   /**
    * Allocate SchemaPartitions
    *
    * @param unassignedSchemaPartitionSlotsMap SchemaPartitionSlots that should be assigned
-   * @return Map<StorageGroupName, SchemaPartitionTable>, the allocating result
+   * @return Map<DatabaseName, SchemaPartitionTable>, the allocating result
    */
   public Map<String, SchemaPartitionTable> allocateSchemaPartition(
       Map<String, List<TSeriesPartitionSlot>> unassignedSchemaPartitionSlotsMap)
       throws NoAvailableRegionGroupException {
-    return genPartitionAllocator().allocateSchemaPartition(unassignedSchemaPartitionSlotsMap);
+    Map<String, SchemaPartitionTable> result = new ConcurrentHashMap<>();
+
+    for (Map.Entry<String, List<TSeriesPartitionSlot>> slotsMapEntry :
+        unassignedSchemaPartitionSlotsMap.entrySet()) {
+      final String database = slotsMapEntry.getKey();
+      final List<TSeriesPartitionSlot> unassignedPartitionSlots = slotsMapEntry.getValue();
+
+      // Filter available SchemaRegionGroups and
+      // sort them by the number of allocated SchemaPartitions
+      BalanceTreeMap<TConsensusGroupId, Integer> counter = new BalanceTreeMap<>();
+      List<Pair<Long, TConsensusGroupId>> regionSlotsCounter =
+          getPartitionManager()
+              .getSortedRegionGroupSlotsCounter(database, TConsensusGroupType.SchemaRegion);
+      for (Pair<Long, TConsensusGroupId> pair : regionSlotsCounter) {
+        counter.put(pair.getRight(), pair.getLeft().intValue());
+      }
+
+      // Enumerate SeriesPartitionSlot
+      Map<TSeriesPartitionSlot, TConsensusGroupId> schemaPartitionMap = new ConcurrentHashMap<>();
+      for (TSeriesPartitionSlot seriesPartitionSlot : unassignedPartitionSlots) {
+        // Greedy allocation: allocate the unassigned SchemaPartition to
+        // the RegionGroup whose allocated SchemaPartitions is the least
+        TConsensusGroupId consensusGroupId = counter.getKeyWithMinValue();
+        schemaPartitionMap.put(seriesPartitionSlot, consensusGroupId);
+        counter.put(consensusGroupId, counter.get(consensusGroupId) + 1);
+      }
+      result.put(database, new SchemaPartitionTable(schemaPartitionMap));
+    }
+
+    return result;
   }
 
   /**
    * Allocate DataPartitions
    *
    * @param unassignedDataPartitionSlotsMap DataPartitionSlots that should be assigned
-   * @return Map<StorageGroupName, DataPartitionTable>, the allocating result
+   * @return Map<DatabaseName, DataPartitionTable>, the allocating result
    */
   public Map<String, DataPartitionTable> allocateDataPartition(
       Map<String, Map<TSeriesPartitionSlot, TTimeSlotList>> unassignedDataPartitionSlotsMap)
       throws NoAvailableRegionGroupException {
-    return genPartitionAllocator().allocateDataPartition(unassignedDataPartitionSlotsMap);
+    Map<String, DataPartitionTable> result = new ConcurrentHashMap<>();
+
+    for (Map.Entry<String, Map<TSeriesPartitionSlot, TTimeSlotList>> slotsMapEntry :
+        unassignedDataPartitionSlotsMap.entrySet()) {
+      final String database = slotsMapEntry.getKey();
+      final Map<TSeriesPartitionSlot, TTimeSlotList> unassignedPartitionSlotsMap =
+          slotsMapEntry.getValue();
+
+      // Filter available DataRegionGroups and
+      // sort them by the number of allocated DataPartitions
+      BalanceTreeMap<TConsensusGroupId, Integer> counter = new BalanceTreeMap<>();
+      List<Pair<Long, TConsensusGroupId>> regionSlotsCounter =
+          getPartitionManager()
+              .getSortedRegionGroupSlotsCounter(database, TConsensusGroupType.DataRegion);
+      for (Pair<Long, TConsensusGroupId> pair : regionSlotsCounter) {
+        counter.put(pair.getRight(), pair.getLeft().intValue());
+      }
+
+      DataPartitionPolicyTable allotTable = dataAllotTableMap.get(database);
+      allotTable.acquireLock();
+      DataPartitionTable dataPartitionTable = new DataPartitionTable();
+
+      try {
+        // Enumerate SeriesPartitionSlot
+        for (Map.Entry<TSeriesPartitionSlot, TTimeSlotList> seriesPartitionEntry :
+            unassignedPartitionSlotsMap.entrySet()) {
+          SeriesPartitionTable seriesPartitionTable = new SeriesPartitionTable();
+
+          // Enumerate TimePartitionSlot in ascending order
+          TSeriesPartitionSlot seriesPartitionSlot = seriesPartitionEntry.getKey();
+          List<TTimePartitionSlot> timePartitionSlots =
+              seriesPartitionEntry.getValue().getTimePartitionSlots();
+          timePartitionSlots.sort(Comparator.comparingLong(TTimePartitionSlot::getStartTime));
+
+          for (TTimePartitionSlot timePartitionSlot : timePartitionSlots) {
+
+            // 1. The historical DataPartition will try to inherit successor DataPartition first
+            TConsensusGroupId successor =
+                getPartitionManager()
+                    .getSuccessorDataPartition(database, seriesPartitionSlot, timePartitionSlot);
+            if (successor != null && counter.containsKey(successor)) {
+              seriesPartitionTable.putDataPartition(timePartitionSlot, successor);
+              counter.put(successor, counter.get(successor) + 1);
+              continue;
+            }
+
+            // 2. Assign DataPartition base on the DataAllotTable
+            TConsensusGroupId allotGroupId =
+                allotTable.getRegionGroupIdOrActivateIfNecessary(seriesPartitionSlot);
+            if (counter.containsKey(allotGroupId)) {
+              seriesPartitionTable.putDataPartition(timePartitionSlot, allotGroupId);
+              counter.put(allotGroupId, counter.get(allotGroupId) + 1);
+              continue;
+            }
+
+            // 3. The allotDataRegionGroup is unavailable,
+            // try to inherit predecessor DataPartition
+            TConsensusGroupId predecessor =
+                getPartitionManager()
+                    .getPredecessorDataPartition(database, seriesPartitionSlot, timePartitionSlot);
+            if (predecessor != null && counter.containsKey(predecessor)) {
+              seriesPartitionTable.putDataPartition(timePartitionSlot, predecessor);
+              counter.put(predecessor, counter.get(predecessor) + 1);
+              continue;
+            }
+
+            // 4. Assign the DataPartition to DataRegionGroup with the least DataPartitions
+            // If the above DataRegionGroups are unavailable
+            TConsensusGroupId greedyGroupId = counter.getKeyWithMinValue();
+            seriesPartitionTable.putDataPartition(timePartitionSlot, greedyGroupId);
+            counter.put(greedyGroupId, counter.get(greedyGroupId) + 1);
+          }
+
+          dataPartitionTable
+              .getDataPartitionMap()
+              .put(seriesPartitionEntry.getKey(), seriesPartitionTable);
+        }
+      } finally {
+        allotTable.releaseLock();
+      }
+      result.put(database, dataPartitionTable);
+    }
+
+    return result;
+  }
+
+  /**
+   * Re-balance the DataPartitionPolicyTable.
+   *
+   * @param database Database name
+   */
+  public void reBalanceDataPartitionPolicy(String database) {
+    try {
+      dataAllotTableMap
+          .computeIfAbsent(database, empty -> new DataPartitionPolicyTable())
+          .reBalanceDataPartitionPolicy(
+              getPartitionManager().getAllRegionGroupIds(database, TConsensusGroupType.DataRegion));
+    } catch (DatabaseNotExistsException e) {
+      LOGGER.error("Database {} not exists when updateDataAllotTable", database);
+    }
+  }
+
+  /** Set up the PartitionBalancer when the current ConfigNode becomes leader. */
+  public void setupPartitionBalancer() {
+    dataAllotTableMap.clear();
+    getClusterSchemaManager()
+        .getDatabaseNames()
+        .forEach(
+            database -> {
+              dataAllotTableMap.put(database, new DataPartitionPolicyTable());

Review Comment:
   ```
   DataPartitionPolicyTable dataPartitionPolicyTable = new DataPartitionPolicyTable();
   dataAllotTableMap.put(database, dataPartitionPolicyTable);
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@iotdb.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org