(iotdb) 02/03: Fix error judgement for MainFragmentLocatedRegion and MostlyUsedDataRegion in distributed plan

[caogaofei]

This is an automated email from the ASF dual-hosted git repository.

caogaofei pushed a commit to branch meituan
in repository https://gitbox.apache.org/repos/asf/iotdb.git

commit 66f45341455c94a28c2e677b913df01808d9b93e
Author: Beyyes <cgf1...@foxmail.com>
AuthorDate: Wed Nov 29 14:53:09 2023 +0800

    Fix error judgement for MainFragmentLocatedRegion and MostlyUsedDataRegion 
in distributed plan
---
 .../planner/distribution/ExchangeNodeAdder.java    | 41 ++++++++++++----------
 1 file changed, 23 insertions(+), 18 deletions(-)

diff --git 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/ExchangeNodeAdder.java
 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/ExchangeNodeAdder.java
index 79d9f6beb96..e95aa68720c 100644
--- 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/ExchangeNodeAdder.java
+++ 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/ExchangeNodeAdder.java
@@ -462,24 +462,29 @@ public class ExchangeNodeAdder extends 
PlanVisitor<PlanNode, NodeGroupContext> {
     // Step 2: return the RegionReplicaSet with max node count
     long maxCount = -1;
     TRegionReplicaSet result = DataPartition.NOT_ASSIGNED;
-    for (Map.Entry<TRegionReplicaSet, Long> entry : groupByRegion.entrySet()) {
-      TRegionReplicaSet region = entry.getKey();
-      if (DataPartition.NOT_ASSIGNED.equals(region)) {
-        continue;
-      }
-      if (region.equals(context.queryContext.getMainFragmentLocatedRegion())) {
-        return context.queryContext.getMainFragmentLocatedRegion();
-      }
-      if (region.equals(context.getMostlyUsedDataRegion())) {
-        return region;
-      }
-      long planNodeCount = entry.getValue();
-      if (planNodeCount > maxCount) {
-        maxCount = planNodeCount;
-        result = region;
-      } else if (planNodeCount == maxCount
-          && region.getRegionId().getId() < result.getRegionId().getId()) {
-        result = region;
+    // use MainFragmentLocatedRegion firstly,
+    // if MainFragmentLocatedRegion is not exist, use MostlyUsedDataRegion,
+    // otherwise use region which has most plan nodes.
+    if (context.queryContext.getMainFragmentLocatedRegion() != null
+        && 
groupByRegion.containsKey(context.queryContext.getMainFragmentLocatedRegion())) 
{
+      return context.queryContext.getMainFragmentLocatedRegion();
+    } else if (context.getMostlyUsedDataRegion() != null
+        && groupByRegion.containsKey(context.getMostlyUsedDataRegion())) {
+      return context.getMostlyUsedDataRegion();
+    } else {
+      for (Map.Entry<TRegionReplicaSet, Long> entry : 
groupByRegion.entrySet()) {
+        TRegionReplicaSet region = entry.getKey();
+        if (DataPartition.NOT_ASSIGNED.equals(region)) {
+          continue;
+        }
+        long planNodeCount = entry.getValue();
+        if (planNodeCount > maxCount) {
+          maxCount = planNodeCount;
+          result = region;
+        } else if (planNodeCount == maxCount
+            && region.getRegionId().getId() < result.getRegionId().getId()) {
+          result = region;
+        }
       }
     }
     return result;