Prefer local DataNode when LOAD creates DataRegions

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/partition/GetDataPartitionPlan.java

@@ -34,6 +34,7 @@ public class GetDataPartitionPlan extends ConfigPhysicalReadPlan {
 
   // Map<StorageGroup, Map<TSeriesPartitionSlot, List<TTimePartitionSlot>>>
   protected Map<String, Map<TSeriesPartitionSlot, TTimeSlotList>> partitionSlotsMap;
+  protected int preferredDataNodeId = -1;
 
   public GetDataPartitionPlan(final ConfigPhysicalPlanType configPhysicalPlanType) {
     super(configPhysicalPlanType);
@@ -49,6 +50,14 @@ public Map<String, Map<TSeriesPartitionSlot, TTimeSlotList>> getPartitionSlotsMa
     return partitionSlotsMap;
   }
 
+  public int getPreferredDataNodeId() {
+    return preferredDataNodeId;
+  }
+
+  public void setPreferredDataNodeId(final int preferredDataNodeId) {
+    this.preferredDataNodeId = preferredDataNodeId;
+  }
+
   /**
    * Convert TDataPartitionReq to GetDataPartitionPlan.
    *
@@ -68,11 +77,12 @@ public boolean equals(final Object o) {
       return false;
     }
     final GetDataPartitionPlan that = (GetDataPartitionPlan) o;
-    return partitionSlotsMap.equals(that.partitionSlotsMap);
+    return partitionSlotsMap.equals(that.partitionSlotsMap)
+        && preferredDataNodeId == that.preferredDataNodeId;
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(partitionSlotsMap);
+    return Objects.hash(partitionSlotsMap, preferredDataNodeId);
   }
 }

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/partition/GetOrCreateDataPartitionPlan.java

@@ -35,6 +35,13 @@ public GetOrCreateDataPartitionPlan(
     this.partitionSlotsMap = partitionSlotsMap;
   }
 
+  public GetOrCreateDataPartitionPlan(
+      final Map<String, Map<TSeriesPartitionSlot, TTimeSlotList>> partitionSlotsMap,
+      final int preferredDataNodeId) {
+    this(partitionSlotsMap);
+    this.preferredDataNodeId = preferredDataNodeId;
+  }
+
   /**
    * Convert TDataPartitionReq to GetOrCreateDataPartitionPlan.
    *
@@ -43,6 +50,11 @@ public GetOrCreateDataPartitionPlan(
    */
   public static GetOrCreateDataPartitionPlan convertFromRpcTDataPartitionReq(
       final TDataPartitionReq req) {
-    return new GetOrCreateDataPartitionPlan(new ConcurrentHashMap<>(req.getPartitionSlotsMap()));
+    final GetOrCreateDataPartitionPlan plan =
+        new GetOrCreateDataPartitionPlan(new ConcurrentHashMap<>(req.getPartitionSlotsMap()));
+    if (req.isSetPreferredDataNodeId()) {
+      plan.setPreferredDataNodeId(req.getPreferredDataNodeId());
+    }
+    return plan;
   }
 }

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/LoadManager.java

@@ -126,6 +126,15 @@ public CreateRegionGroupsPlan allocateRegionGroups(
     return regionBalancer.genRegionGroupsAllocationPlan(allotmentMap, consensusGroupType);
   }
 
+  public CreateRegionGroupsPlan allocateRegionGroups(
+      final Map<String, Integer> allotmentMap,
+      final TConsensusGroupType consensusGroupType,
+      final Map<String, Integer> preferredDataNodeMap)
+      throws NotEnoughDataNodeException, DatabaseNotExistsException {
+    return regionBalancer.genRegionGroupsAllocationPlan(
+        allotmentMap, consensusGroupType, preferredDataNodeMap);
+  }
+
   /**
    * Allocate SchemaPartitions.
    *

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/RegionBalancer.java

@@ -22,6 +22,7 @@
 import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
 import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
 import org.apache.iotdb.common.rpc.thrift.TDataNodeConfiguration;
+import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
 import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
 import org.apache.iotdb.commons.cluster.NodeStatus;
 import org.apache.iotdb.confignode.conf.ConfigNodeDescriptor;
@@ -39,6 +40,8 @@
 import org.apache.iotdb.confignode.manager.partition.PartitionManager;
 import org.apache.iotdb.confignode.manager.schema.ClusterSchemaManager;
 
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -82,6 +85,14 @@ public RegionBalancer(IManager configManager) {
   public CreateRegionGroupsPlan genRegionGroupsAllocationPlan(
       final Map<String, Integer> allotmentMap, final TConsensusGroupType consensusGroupType)
       throws NotEnoughDataNodeException, DatabaseNotExistsException {
+    return genRegionGroupsAllocationPlan(allotmentMap, consensusGroupType, null);
+  }
+
+  public CreateRegionGroupsPlan genRegionGroupsAllocationPlan(
+      final Map<String, Integer> allotmentMap,
+      final TConsensusGroupType consensusGroupType,
+      final Map<String, Integer> preferredDataNodeMap)
+      throws NotEnoughDataNodeException, DatabaseNotExistsException {
 
     // Some new RegionGroups will have to occupy unknown DataNodes if the number of online
     // DataNodes is insufficient (Unknown DataNodes are intentionally kept as candidates).
@@ -118,6 +129,12 @@ public CreateRegionGroupsPlan genRegionGroupsAllocationPlan(
       final int allotment = entry.getValue();
       final int replicationFactor =
           getClusterSchemaManager().getReplicationFactor(database, consensusGroupType);
+      final Integer requestPreferredDataNode =
+          preferredDataNodeMap == null ? null : preferredDataNodeMap.get(database);
+      final int preferredDataNodeId =
+          TConsensusGroupType.DataRegion.equals(consensusGroupType)
+              ? (requestPreferredDataNode != null ? requestPreferredDataNode : -1)
+              : -1;
       // Only considering the specified Database when doing allocation
       final List<TRegionReplicaSet> databaseAllocatedRegionGroups =
           getPartitionManager().getAllReplicaSets(database, consensusGroupType);
@@ -144,6 +161,7 @@ public CreateRegionGroupsPlan genRegionGroupsAllocationPlan(
                 replicationFactor,
                 new TConsensusGroupId(
                     consensusGroupType, getPartitionManager().generateNextRegionGroupId()));
+        preferDataNodeIfPossible(newRegionGroup, preferredDataNodeId, availableDataNodeMap);
         createRegionGroupsPlan.addRegionGroup(database, newRegionGroup);
 
         // Mark the new RegionGroup as allocated
@@ -175,6 +193,37 @@ private ProcedureManager getProcedureManager() {
     return configManager.getProcedureManager();
   }
 
+  private static void preferDataNodeIfPossible(
+      final TRegionReplicaSet regionReplicaSet,
+      final int preferredDataNodeId,
+      final Map<Integer, TDataNodeConfiguration> availableDataNodeMap) {
+    if (preferredDataNodeId < 0 || !availableDataNodeMap.containsKey(preferredDataNodeId)) {
+      return;
+    }
+
+    final List<TDataNodeLocation> locations =
+        new ArrayList<>(regionReplicaSet.getDataNodeLocations());
+    if (locations.isEmpty()) {
+      return;
+    }
+
+    int preferredIndex = -1;
+    for (int i = 0; i < locations.size(); i++) {
+      if (locations.get(i).getDataNodeId() == preferredDataNodeId) {
+        preferredIndex = i;
+        break;
+      }
+    }
+
+    if (preferredIndex > 0) {
+      Collections.swap(locations, 0, preferredIndex);
+    } else if (preferredIndex < 0) {
+      locations.set(
+          locations.size() - 1, availableDataNodeMap.get(preferredDataNodeId).getLocation());
+    }
+    regionReplicaSet.setDataNodeLocations(locations);
+  }
+
   public enum RegionGroupAllocatePolicy {
     GREEDY,
     GCR,

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/partition/PartitionManager.java

@@ -446,7 +446,9 @@ public DataPartitionResp getOrCreateDataPartition(final GetOrCreateDataPartition
                   storageGroup, unassignedDataPartitionSlots.size()));
       TSStatus status =
           extendRegionGroupIfNecessary(
-              unassignedDataPartitionSlotsCountMap, TConsensusGroupType.DataRegion);
+              unassignedDataPartitionSlotsCountMap,
+              TConsensusGroupType.DataRegion,
+              req.getPreferredDataNodeId());
       if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
         // Return an error code if Region extension failed
         resp.setStatus(status);
@@ -607,6 +609,13 @@ private TSStatus consensusWritePartitionResult(ConfigPhysicalPlan plan) {
   private TSStatus extendRegionGroupIfNecessary(
       final Map<String, Integer> unassignedPartitionSlotsCountMap,
       final TConsensusGroupType consensusGroupType) {
+    return extendRegionGroupIfNecessary(unassignedPartitionSlotsCountMap, consensusGroupType, -1);
+  }
+
+  private TSStatus extendRegionGroupIfNecessary(
+      final Map<String, Integer> unassignedPartitionSlotsCountMap,
+      final TConsensusGroupType consensusGroupType,
+      final int preferredDataNodeId) {
 
     final TSStatus result = new TSStatus();
 
@@ -615,21 +624,21 @@ private TSStatus extendRegionGroupIfNecessary(
         switch (CONF.getSchemaRegionGroupExtensionPolicy()) {
           case CUSTOM:
             return customExtendRegionGroupIfNecessary(
-                unassignedPartitionSlotsCountMap, consensusGroupType);
+                unassignedPartitionSlotsCountMap, consensusGroupType, preferredDataNodeId);
           case AUTO:
           default:
             return autoExtendRegionGroupIfNecessary(
-                unassignedPartitionSlotsCountMap, consensusGroupType);
+                unassignedPartitionSlotsCountMap, consensusGroupType, preferredDataNodeId);
         }
       } else {
         switch (CONF.getDataRegionGroupExtensionPolicy()) {
           case CUSTOM:
             return customExtendRegionGroupIfNecessary(
-                unassignedPartitionSlotsCountMap, consensusGroupType);
+                unassignedPartitionSlotsCountMap, consensusGroupType, preferredDataNodeId);
           case AUTO:
           default:
             return autoExtendRegionGroupIfNecessary(
-                unassignedPartitionSlotsCountMap, consensusGroupType);
+                unassignedPartitionSlotsCountMap, consensusGroupType, preferredDataNodeId);
         }
       }
     } catch (NotEnoughDataNodeException e) {
@@ -647,7 +656,8 @@ private TSStatus extendRegionGroupIfNecessary(
 
   private TSStatus customExtendRegionGroupIfNecessary(
       final Map<String, Integer> unassignedPartitionSlotsCountMap,
-      final TConsensusGroupType consensusGroupType)
+      final TConsensusGroupType consensusGroupType,
+      final int preferredDataNodeId)
       throws DatabaseNotExistsException, NotEnoughDataNodeException {
 
     // Map<Database, Region allotment>
@@ -666,12 +676,13 @@ private TSStatus customExtendRegionGroupIfNecessary(
       }
     }
 
-    return generateAndAllocateRegionGroups(allotmentMap, consensusGroupType);
+    return generateAndAllocateRegionGroups(allotmentMap, consensusGroupType, preferredDataNodeId);
   }
 
   private TSStatus autoExtendRegionGroupIfNecessary(
       final Map<String, Integer> unassignedPartitionSlotsCountMap,
-      final TConsensusGroupType consensusGroupType)
+      final TConsensusGroupType consensusGroupType,
+      final int preferredDataNodeId)
       throws NotEnoughDataNodeException, DatabaseNotExistsException {
 
     // Map<Database, Region allotment>
@@ -735,15 +746,24 @@ private TSStatus autoExtendRegionGroupIfNecessary(
       }
     }
 
-    return generateAndAllocateRegionGroups(allotmentMap, consensusGroupType);
+    return generateAndAllocateRegionGroups(allotmentMap, consensusGroupType, preferredDataNodeId);
   }
 
   private TSStatus generateAndAllocateRegionGroups(
-      final Map<String, Integer> allotmentMap, final TConsensusGroupType consensusGroupType)
+      final Map<String, Integer> allotmentMap,
+      final TConsensusGroupType consensusGroupType,
+      final int preferredDataNodeId)
       throws NotEnoughDataNodeException, DatabaseNotExistsException {
     if (!allotmentMap.isEmpty()) {
+      final Map<String, Integer> preferredDataNodeMap = new ConcurrentHashMap<>();
+      if (preferredDataNodeId >= 0) {
+        allotmentMap
+            .keySet()
+            .forEach(database -> preferredDataNodeMap.put(database, preferredDataNodeId));
+      }
       final CreateRegionGroupsPlan createRegionGroupsPlan =
-          getLoadManager().allocateRegionGroups(allotmentMap, consensusGroupType);
+          getLoadManager()
+              .allocateRegionGroups(allotmentMap, consensusGroupType, preferredDataNodeMap);
       LOGGER.info(ManagerMessages.CREATEREGIONGROUPS_STARTING_TO_CREATE_THE_FOLLOWING_REGIONGROUPS);
       createRegionGroupsPlan.planLog(LOGGER);
       return getProcedureManager().createRegionGroups(consensusGroupType, createRegionGroupsPlan);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java

@@ -1156,6 +1156,9 @@ public class IoTDBConfig {
   /** Load related */
   private double maxAllocateMemoryRatioForLoad = 0.8;
 
+  /** Hidden hot-reload switch that prefers placing LOAD-created DataRegions on this DataNode. */
+  private volatile boolean loadTsFilePreferLocalNode = true;
+
   private int loadTsFileAnalyzeSchemaBatchReadTimeSeriesMetadataCount = 4096;
   private int loadTsFileAnalyzeSchemaBatchFlushTimeSeriesNumber = 4096;
   private int loadTsFileAnalyzeSchemaBatchFlushTableDeviceNumber = 4096; // For table model
@@ -4289,6 +4292,14 @@ public void setLoadTsFileRetryCountOnRegionChange(int loadTsFileRetryCountOnRegi
     this.loadTsFileRetryCountOnRegionChange = loadTsFileRetryCountOnRegionChange;
   }
 
+  public boolean isLoadTsFilePreferLocalNode() {
+    return loadTsFilePreferLocalNode;
+  }
+
+  public void setLoadTsFilePreferLocalNode(final boolean loadTsFilePreferLocalNode) {
+    this.loadTsFilePreferLocalNode = loadTsFilePreferLocalNode;
+  }
+
   public double getLoadWriteThroughputBytesPerSecond() {
     return loadWriteThroughputBytesPerSecond;
   }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java

@@ -2642,6 +2642,11 @@ private void loadLoadTsFileProps(TrimProperties properties) {
             properties.getProperty(
                 "load_tsfile_source_path_check_enable",
                 Boolean.toString(conf.isLoadTsFileSourcePathCheckEnabled()))));
+    conf.setLoadTsFilePreferLocalNode(
+        Boolean.parseBoolean(
+            properties.getProperty(
+                "load_tsfile_prefer_local_node",
+                Boolean.toString(conf.isLoadTsFilePreferLocalNode()))));
 
     conf.setLoadTabletConversionThresholdBytes(
         Long.parseLong(
@@ -2791,6 +2796,11 @@ private void loadLoadTsFileHotModifiedProp(TrimProperties properties) throws IOE
             properties.getProperty(
                 "load_tsfile_source_path_check_enable",
                 Boolean.toString(conf.isLoadTsFileSourcePathCheckEnabled()))));
+    conf.setLoadTsFilePreferLocalNode(
+        Boolean.parseBoolean(
+            properties.getProperty(
+                "load_tsfile_prefer_local_node",
+                Boolean.toString(conf.isLoadTsFilePreferLocalNode()))));
 
     conf.setLoadActiveListeningEnable(
         Boolean.parseBoolean(

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ClusterPartitionFetcher.java

@@ -75,6 +75,7 @@
 public class ClusterPartitionFetcher implements IPartitionFetcher {
 
   private static final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig();
+  private static final ThreadLocal<Integer> LOAD_PREFERRED_DATA_NODE_ID = new ThreadLocal<>();
 
   private final SeriesPartitionExecutor partitionExecutor;
 
@@ -94,6 +95,23 @@ public static ClusterPartitionFetcher getInstance() {
     return ClusterPartitionFetcherHolder.INSTANCE;
   }
 
+  public static void setLoadPreferredDataNodeId(final int preferredDataNodeId) {
+    if (preferredDataNodeId >= 0) {
+      LOAD_PREFERRED_DATA_NODE_ID.set(preferredDataNodeId);
+    } else {
+      LOAD_PREFERRED_DATA_NODE_ID.remove();
+    }
+  }
+
+  public static void clearLoadPreferredDataNodeId() {
+    LOAD_PREFERRED_DATA_NODE_ID.remove();
+  }
+
+  private static int getLoadPreferredDataNodeId() {
+    final Integer preferredDataNodeId = LOAD_PREFERRED_DATA_NODE_ID.get();
+    return preferredDataNodeId == null ? -1 : preferredDataNodeId;
+  }
+
   private ClusterPartitionFetcher() {
     this.partitionExecutor =
         SeriesPartitionExecutor.getSeriesPartitionExecutor(
@@ -319,6 +337,15 @@ public DataPartition getOrCreateDataPartition(
   @Override
   public DataPartition getOrCreateDataPartition(
       final List<DataPartitionQueryParam> dataPartitionQueryParams, final String userName) {
+    return getOrCreateDataPartition(
+        dataPartitionQueryParams, userName, getLoadPreferredDataNodeId());
+  }
+
+  @Override
+  public DataPartition getOrCreateDataPartition(
+      final List<DataPartitionQueryParam> dataPartitionQueryParams,
+      final String userName,
+      final int preferredDataNodeId) {
     final Map<String, List<DataPartitionQueryParam>> splitDataPartitionQueryParams =
         splitDataPartitionQueryParam(
             dataPartitionQueryParams, config.isAutoCreateSchemaEnabled(), userName);
@@ -330,6 +357,9 @@ public DataPartition getOrCreateDataPartition(
     try (final ConfigNodeClient client =
         configNodeClientManager.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
       final TDataPartitionReq req = constructDataPartitionReq(splitDataPartitionQueryParams);
+      if (preferredDataNodeId >= 0) {
+        req.setPreferredDataNodeId(preferredDataNodeId);
+      }
       final TDataPartitionTableResp dataPartitionTableResp =
           client.getOrCreateDataPartitionTable(req);
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IPartitionFetcher.java

@@ -100,6 +100,13 @@ DataPartition getOrCreateDataPartition(
   DataPartition getOrCreateDataPartition(
       final List<DataPartitionQueryParam> dataPartitionQueryParams, final String userName);
 
+  default DataPartition getOrCreateDataPartition(
+      final List<DataPartitionQueryParam> dataPartitionQueryParams,
+      final String userName,
+      final int preferredDataNodeId) {
+    return getOrCreateDataPartition(dataPartitionQueryParams, userName);
+  }
+
   /** Get schema partition and matched nodes according to path pattern tree. */
   default SchemaNodeManagementPartition getSchemaNodeManagementPartition(
       PathPatternTree patternTree, PathPatternTree scope, boolean needAuditDB) {