Fix forced Exchange branch isolation in distributed align-by-device planning.

This preserves forced InnerTimeJoin-related exchanges from being merged into shared sinks and adds an integration test for wildcard align-by-device aggregation across time partitions.

Author: Wei-hao-Li

integration-test/src/test/java/org/apache/iotdb/db/it/aggregation/IoTDBAlignByDeviceWildcardIT.java

@@ -0,0 +1,94 @@
+/*
+ * 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.db.it.aggregation;
+
+import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.LocalStandaloneIT;
+
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import static org.apache.iotdb.db.it.utils.TestUtils.prepareData;
+import static org.apache.iotdb.db.it.utils.TestUtils.resultSetEqualTest;
+
+@RunWith(IoTDBTestRunner.class)
+@Category({LocalStandaloneIT.class})
+public class IoTDBAlignByDeviceWildcardIT {
+
+  private static final String[] SQL_LIST =
+      new String[] {
+        "CREATE DATABASE root.min",
+        "CREATE TIMESERIES root.min.d1.s1 WITH DATATYPE=INT32, ENCODING=PLAIN",
+        "CREATE TIMESERIES root.min.d1.s2 WITH DATATYPE=INT32, ENCODING=PLAIN",
+        "CREATE TIMESERIES root.min.d2.s1 WITH DATATYPE=INT32, ENCODING=PLAIN",
+        "CREATE TIMESERIES root.min.d2.s2 WITH DATATYPE=INT32, ENCODING=PLAIN",
+        "INSERT INTO root.min.d1(time, s1, s2) VALUES(1, 1, 1)",
+        "INSERT INTO root.min.d1(time, s1, s2) VALUES(2, 1, 1)",
+        "FLUSH",
+        "INSERT INTO root.min.d2(time, s1, s2) VALUES(5, 1, 1)",
+        "FLUSH"
+      };
+
+  @BeforeClass
+  public static void setUp() throws Exception {
+    EnvFactory.getEnv()
+        .getConfig()
+        .getCommonConfig()
+        .setDefaultDataRegionGroupNumPerDatabase(1)
+        .setTimePartitionInterval(1)
+        .setEnableSeqSpaceCompaction(false)
+        .setEnableUnseqSpaceCompaction(false)
+        .setEnableCrossSpaceCompaction(false);
+    EnvFactory.getEnv().initClusterEnvironment();
+    prepareData(SQL_LIST);
+  }
+
+  @AfterClass
+  public static void tearDown() throws Exception {
+    // EnvFactory.getEnv().cleanClusterEnvironment();
+  }
+
+  @Test
+  public void testWildcardAlignByDeviceWithTimePartitionSplit() {
+    String sql =
+        "SELECT count(s1) FROM root.min.** "
+            + "WHERE s2 is not null and s1 is not null "
+            + "GROUP BY([1, 6), 1ms) ALIGN BY DEVICE";
+    String[] expectedHeader = new String[] {"Time", "Device", "count(s1)"};
+    String[] expectedRows =
+        new String[] {
+          "1,root.min.d1,1,",
+          "2,root.min.d1,1,",
+          "3,root.min.d1,0,",
+          "4,root.min.d1,0,",
+          "5,root.min.d1,0,",
+          "1,root.min.d2,0,",
+          "2,root.min.d2,0,",
+          "3,root.min.d2,0,",
+          "4,root.min.d2,0,",
+          "5,root.min.d2,1,",
+        };
+    resultSetEqualTest(sql, expectedHeader, expectedRows);
+  }
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/DistributionPlanner.java

@@ -145,13 +145,23 @@ private void adjustUpStreamHelper(
         ExchangeNode exchangeNode = (ExchangeNode) child;
         TRegionReplicaSet regionOfChild =
             context.getNodeDistribution(exchangeNode.getChild().getPlanNodeId()).getRegion();
-        MultiChildrenSinkNode newChild =
-            memo.computeIfAbsent(
-                regionOfChild,
-                tRegionReplicaSet ->
-                    needShuffleSinkNode
-                        ? new ShuffleSinkNode(context.queryContext.getQueryId().genPlanNodeId())
-                        : new IdentitySinkNode(context.queryContext.getQueryId().genPlanNodeId()));
+        MultiChildrenSinkNode newChild;
+        if (exchangeNode.isForcedExchange()) {
+          // Keep forced exchange branch isolated: do not merge into shared sink memo.
+          newChild =
+              needShuffleSinkNode
+                  ? new ShuffleSinkNode(context.queryContext.getQueryId().genPlanNodeId())
+                  : new IdentitySinkNode(context.queryContext.getQueryId().genPlanNodeId());
+        } else {
+          newChild =
+              memo.computeIfAbsent(
+                  regionOfChild,
+                  tRegionReplicaSet ->
+                      needShuffleSinkNode
+                          ? new ShuffleSinkNode(context.queryContext.getQueryId().genPlanNodeId())
+                          : new IdentitySinkNode(
+                              context.queryContext.getQueryId().genPlanNodeId()));
+        }
         newChild.addChild(exchangeNode.getChild());
         newChild.addDownStreamChannelLocation(
             new DownStreamChannelLocation(exchangeNode.getPlanNodeId().toString()));

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/ExchangeNodeAdder.java

@@ -69,17 +69,17 @@
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.SeriesScanNode;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.SourceNode;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 
-import static com.google.common.collect.ImmutableList.toImmutableList;
-
 public class ExchangeNodeAdder extends PlanVisitor<PlanNode, NodeGroupContext> {
 
   private final Analysis analysis;
+  private boolean containsInnerTimeJoinInCurrentSubtree = false;
 
   public ExchangeNodeAdder(Analysis analysis) {
     this.analysis = analysis;
@@ -92,10 +92,7 @@ public PlanNode visitPlan(PlanNode node, NodeGroupContext context) {
       return node;
     }
     // Visit all the children of current node
-    List<PlanNode> children =
-        node.getChildren().stream()
-            .map(child -> child.accept(this, context))
-            .collect(toImmutableList());
+    List<PlanNode> children = visitChildrenAndRecordInnerTimeJoin(node.getChildren(), context);
 
     // Calculate the node distribution info according to its children
 
@@ -215,7 +212,13 @@ private PlanNode processNoChildSourceNode(SourceNode node, NodeGroupContext cont
 
   @Override
   public PlanNode visitDeviceView(DeviceViewNode node, NodeGroupContext context) {
-    return processMultiChildNode(node, context);
+    List<PlanNode> visitedChildren =
+        visitChildrenAndRecordInnerTimeJoin(node.getChildren(), context);
+    // Force Exchange for multi-child DeviceView if any child subtree contains InnerTimeJoin.
+    if (node.getChildren().size() > 1 && containsInnerTimeJoinInCurrentSubtree) {
+      return processMultiChildNodeWithForcedExchange(node, context, visitedChildren);
+    }
+    return processMultiChildNode(node, context, visitedChildren);
   }
 
   @Override
@@ -236,7 +239,13 @@ public PlanNode visitSingleDeviceView(SingleDeviceViewNode node, NodeGroupContex
 
   @Override
   public PlanNode visitMergeSort(MergeSortNode node, NodeGroupContext context) {
-    return processMultiChildNode(node, context);
+    List<PlanNode> visitedChildren =
+        visitChildrenAndRecordInnerTimeJoin(node.getChildren(), context);
+    // Force Exchange if any child subtree contains InnerTimeJoin.
+    if (containsInnerTimeJoinInCurrentSubtree) {
+      return processMultiChildNodeWithForcedExchange(node, context, visitedChildren);
+    }
+    return processMultiChildNode(node, context, visitedChildren);
   }
 
   @Override
@@ -428,11 +437,14 @@ private PlanNode processMultiChildNode(MultiChildProcessNode node, NodeGroupCont
       return processMultiChildNodeByLocation(node, context);
     }
 
-    MultiChildProcessNode newNode = (MultiChildProcessNode) node.clone();
     List<PlanNode> visitedChildren =
-        node.getChildren().stream()
-            .map(child -> visit(child, context))
-            .collect(Collectors.toList());
+        visitChildrenAndRecordInnerTimeJoin(node.getChildren(), context);
+    return processMultiChildNode(node, context, visitedChildren);
+  }
+
+  private PlanNode processMultiChildNode(
+      MultiChildProcessNode node, NodeGroupContext context, List<PlanNode> visitedChildren) {
+    MultiChildProcessNode newNode = (MultiChildProcessNode) node.clone();
 
     // DataRegion in which node locates
     TRegionReplicaSet dataRegion;
@@ -546,13 +558,47 @@ private PlanNode processTopKNode(
   }
 
   private ExchangeNode genExchangeNode(NodeGroupContext context, PlanNode child) {
+    return genExchangeNode(context, child, false);
+  }
+
+  private ExchangeNode genExchangeNode(
+      NodeGroupContext context, PlanNode child, boolean forcedExchange) {
     ExchangeNode exchangeNode = new ExchangeNode(context.queryContext.getQueryId().genPlanNodeId());
     exchangeNode.setChild(child);
     exchangeNode.setOutputColumnNames(child.getOutputColumnNames());
+    exchangeNode.setForcedExchange(forcedExchange);
     context.hasExchangeNode = true;
     return exchangeNode;
   }
 
+  private PlanNode processMultiChildNodeWithForcedExchange(
+      MultiChildProcessNode node, NodeGroupContext context, List<PlanNode> visitedChildren) {
+    MultiChildProcessNode newNode = (MultiChildProcessNode) node.clone();
+    for (PlanNode child : visitedChildren) {
+      newNode.addChild(genExchangeNode(context, child, true));
+    }
+    context.putNodeDistribution(
+        newNode.getPlanNodeId(),
+        new NodeDistribution(
+            NodeDistributionType.SAME_WITH_SOME_CHILD, context.getMostlyUsedDataRegion()));
+    return newNode;
+  }
+
+  private List<PlanNode> visitChildrenAndRecordInnerTimeJoin(
+      List<PlanNode> children, NodeGroupContext context) {
+    List<PlanNode> result = new ArrayList<>(children.size());
+    boolean originalTimeJoin = containsInnerTimeJoinInCurrentSubtree;
+    boolean hasInnerTimeJoinInChildren = false;
+    for (PlanNode child : children) {
+      containsInnerTimeJoinInCurrentSubtree = false;
+      PlanNode visitedChild = visit(child, context);
+      hasInnerTimeJoinInChildren |= containsInnerTimeJoinInCurrentSubtree;
+      result.add(visitedChild);
+    }
+    containsInnerTimeJoinInCurrentSubtree = originalTimeJoin || hasInnerTimeJoinInChildren;
+    return result;
+  }
+
   @Override
   public PlanNode visitSlidingWindowAggregation(
       SlidingWindowAggregationNode node, NodeGroupContext context) {
@@ -584,9 +630,9 @@ private TRegionReplicaSet calculateDataRegionByChildren(
                       if (region == null
                           && context.getNodeDistribution(child.getPlanNodeId()).getType()
                               == NodeDistributionType.SAME_WITH_ALL_CHILDREN) {
-                        return calculateSchemaRegionByChildren(child.getChildren(), context);
+                        region = calculateSchemaRegionByChildren(child.getChildren(), context);
                       }
-                      return region;
+                      return region == null ? DataPartition.NOT_ASSIGNED : region;
                     },
                     Collectors.counting()));
 
@@ -644,6 +690,13 @@ private boolean nodeDistributionIsSame(List<PlanNode> children, NodeGroupContext
   }
 
   public PlanNode visit(PlanNode node, NodeGroupContext context) {
-    return node.accept(this, context);
+    boolean originalTimeJoin = containsInnerTimeJoinInCurrentSubtree;
+    containsInnerTimeJoinInCurrentSubtree = false;
+    PlanNode visitedNode = node.accept(this, context);
+    containsInnerTimeJoinInCurrentSubtree =
+        originalTimeJoin
+            || containsInnerTimeJoinInCurrentSubtree
+            || node instanceof InnerTimeJoinNode;
+    return visitedNode;
   }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/process/ExchangeNode.java

@@ -48,6 +48,9 @@ public class ExchangeNode extends SingleChildProcessNode {
   /** Exchange needs to know which child of IdentitySinkNode/ShuffleSinkNode it matches */
   private int indexOfUpstreamSinkHandle = 0;
 
+  /** Planner-only flag: this exchange is forced and should keep independent upstream sink. */
+  private transient boolean forcedExchange = false;
+
   public ExchangeNode(PlanNodeId id) {
     super(id);
   }
@@ -72,6 +75,7 @@ public PlanNode clone() {
     ExchangeNode node = new ExchangeNode(getPlanNodeId());
     node.setOutputColumnNames(outputColumnNames);
     node.setIndexOfUpstreamSinkHandle(indexOfUpstreamSinkHandle);
+    node.setForcedExchange(forcedExchange);
     return node;
   }
 
@@ -162,6 +166,14 @@ public void setIndexOfUpstreamSinkHandle(int indexOfUpstreamSinkHandle) {
     this.indexOfUpstreamSinkHandle = indexOfUpstreamSinkHandle;
   }
 
+  public boolean isForcedExchange() {
+    return forcedExchange;
+  }
+
+  public void setForcedExchange(boolean forcedExchange) {
+    this.forcedExchange = forcedExchange;
+  }
+
   public TEndPoint getUpstreamEndpoint() {
     return upstreamEndpoint;
   }
@@ -188,11 +200,13 @@ public boolean equals(Object o) {
     ExchangeNode that = (ExchangeNode) o;
     return Objects.equals(upstreamEndpoint, that.upstreamEndpoint)
         && Objects.equals(upstreamInstanceId, that.upstreamInstanceId)
-        && Objects.equals(upstreamPlanNodeId, that.upstreamPlanNodeId);
+        && Objects.equals(upstreamPlanNodeId, that.upstreamPlanNodeId)
+        && forcedExchange == that.forcedExchange;
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(super.hashCode(), upstreamEndpoint, upstreamInstanceId, upstreamPlanNodeId);
+    return Objects.hash(
+        super.hashCode(), upstreamEndpoint, upstreamInstanceId, upstreamPlanNodeId, forcedExchange);
   }
 }