Limit shared exchange budgets and ready Collect to table fragments

Author: ColinLeeo

integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java

@@ -34,7 +34,9 @@
 import java.sql.ResultSet;
 import java.sql.Statement;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.assertUnorderedQueryResult;
 import static org.junit.Assert.assertEquals;
@@ -109,7 +111,7 @@ public void testUnorderedUnionPreservesEveryDuplicateOnEachNode() throws Excepti
   }
 
   @Test
-  public void testSharedExchangePathsForTreeAndTreeToTable() throws Exception {
+  public void testLegacyTreeAndSharedTreeToTablePaths() throws Exception {
     try (Connection connection = EnvFactory.getEnv().getConnection();
         Statement statement = connection.createStatement()) {
       statement.execute("CREATE DATABASE root.collect_tree");
@@ -124,10 +126,25 @@ public void testSharedExchangePathsForTreeAndTreeToTable() throws Exception {
       try (ResultSet result =
           statement.executeQuery("SELECT value FROM root.collect_tree.*.* ALIGN BY DEVICE")) {
         int rows = 0;
+        Set<String> seenDevices = new HashSet<>();
+        String currentDevice = null;
+        int deviceRows = 0;
         while (result.next()) {
+          String nextDevice = result.getString("Device");
+          if (!nextDevice.equals(currentDevice)) {
+            if (currentDevice != null) {
+              assertEquals(ROWS, deviceRows);
+            }
+            assertTrue("Tree results must retain contiguous devices", seenDevices.add(nextDevice));
+            currentDevice = nextDevice;
+            deviceRows = 0;
+          }
+          assertEquals(++deviceRows, result.getLong("Time"));
           assertEquals(result.getLong("Time"), result.getLong("value"));
           rows++;
         }
+        assertEquals(ROWS, deviceRows);
+        assertEquals(DEVICES, seenDevices.size());
         assertEquals(DEVICES * ROWS, rows);
       }
     }
@@ -300,7 +317,7 @@ public void testOrderByAfterGlobalWindowAndConstantFill() throws Exception {
   }
 
   @Test
-  public void testTopKAfterGlobalWindowConsumesUnorderedInput() throws Exception {
+  public void testExcludedTopKPlanRetainsQueryResults() throws Exception {
     try (Connection connection = EnvFactory.getEnv().getTableConnection();
         Statement statement = connection.createStatement()) {
       statement.execute("USE " + DATABASE);

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java

@@ -39,7 +39,9 @@ public class CollectOperator implements ProcessOperator {
 
   private final CommonOperatorContext operatorContext;
   protected final List<Operator> children;
-  private final AnyChildBlocked anyChildBlocked;
+  private AnyChildBlocked anyChildBlocked;
+  private boolean inited;
+  protected int currentIndex;
   private ListenableFuture<Integer> selectionFuture;
   private ListenableFuture<Integer> exposedFuture;
   private int remainingChildren;
@@ -48,17 +50,36 @@ public class CollectOperator implements ProcessOperator {
   public CollectOperator(CommonOperatorContext operatorContext, List<Operator> children) {
     this.operatorContext = operatorContext;
     this.children = children;
-    this.anyChildBlocked = new AnyChildBlocked(children);
     this.remainingChildren = children.size();
   }
 
+  /**
+   * Enable only after the complete table FI has passed eligibility checks, before drivers start.
+   */
+  public void enableReadyInputConsumption() {
+    if (anyChildBlocked == null) {
+      anyChildBlocked = new AnyChildBlocked(children);
+    }
+  }
+
   @Override
   public boolean hasNext() {
-    return !closed && remainingChildren > 0;
+    return anyChildBlocked == null
+        ? currentIndex < children.size()
+        : !closed && remainingChildren > 0;
   }
 
   @Override
   public TsBlock next() throws Exception {
+    if (anyChildBlocked == null) {
+      if (children.get(currentIndex).hasNextWithTimer()) {
+        TsBlock block = children.get(currentIndex).nextWithTimer();
+        return block == null ? null : transformBlock(currentIndex, block);
+      }
+      closeCurrentChild(currentIndex);
+      currentIndex++;
+      return null;
+    }
     anyChildBlocked.checkFailure();
     if (!hasNext()) {
       return null;
@@ -98,6 +119,11 @@ protected TsBlock transformBlock(int childId, TsBlock block) {
   }
 
   protected void closeCurrentChild(int index) throws Exception {
+    if (anyChildBlocked == null) {
+      children.get(index).close();
+      children.set(index, null);
+      return;
+    }
     Operator child = children.get(index);
     if (child == null) {
       return;
@@ -111,6 +137,15 @@ protected void closeCurrentChild(int index) throws Exception {
 
   @Override
   public ListenableFuture<?> isBlocked() {
+    if (anyChildBlocked == null) {
+      if (!inited) {
+        inited = true;
+        for (Operator child : children) {
+          child.isBlocked();
+        }
+      }
+      return currentIndex >= children.size() ? NOT_BLOCKED : children.get(currentIndex).isBlocked();
+    }
     if (!hasNext()) {
       return NOT_BLOCKED;
     }
@@ -135,6 +170,14 @@ public CommonOperatorContext getOperatorContext() {
 
   @Override
   public void close() throws Exception {
+    if (anyChildBlocked == null) {
+      for (int i = currentIndex; i < children.size(); i++) {
+        if (children.get(i) != null) {
+          closeCurrentChild(i);
+        }
+      }
+      return;
+    }
     if (closed) {
       return;
     }
@@ -161,6 +204,13 @@ public void close() throws Exception {
 
   @Override
   public long calculateMaxPeekMemory() {
+    if (anyChildBlocked == null) {
+      long maxPeekMemory = 0;
+      for (Operator child : children) {
+        maxPeekMemory = Math.max(maxPeekMemory, child.calculateMaxPeekMemoryWithCounter());
+      }
+      return maxPeekMemory;
+    }
     long retained = calculateRetainedSizeAfterCallingNext();
     long peak = retained;
     for (Operator child : children) {
@@ -189,6 +239,9 @@ public long calculateMaxReturnSize() {
 
   @Override
   public long calculateRetainedSizeAfterCallingNext() {
+    if (anyChildBlocked == null) {
+      return 0;
+    }
     long retained = 0;
     for (Operator child : children) {
       if (child != null) {
@@ -206,7 +259,7 @@ public List<Operator> getChildren() {
   @Override
   public long ramBytesUsed() {
     return INSTANCE_SIZE
-        + anyChildBlocked.ramBytesUsed()
+        + (anyChildBlocked == null ? 0 : anyChildBlocked.ramBytesUsed())
         + children.stream()
             .mapToLong(MemoryEstimationHelper::getEstimatedSizeOfAccountableObject)
             .sum()

iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java

@@ -71,7 +71,10 @@ private static TsBlock block(int... values) {
   }
 
   private static CollectOperator collect(Operator... children) {
-    return new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(children)));
+    CollectOperator collect =
+        new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(children)));
+    collect.enableReadyInputConsumption();
+    return collect;
   }
 
   @Test
@@ -100,6 +103,7 @@ public void testThousandChildrenAndOneListenerPerPendingPeriod() throws Exceptio
         children.add(new Child(block(i)));
       }
       CollectOperator collect = new CollectOperator(new TestContext(), children);
+      collect.enableReadyInputConsumption();
       Child ready = (Child) children.get(count - 1);
       for (int round = 0; round < 5; round++) {
         ListenableFuture<?> wait = collect.isBlocked();
@@ -133,6 +137,7 @@ public void testAllReadyChildrenHaveLinearSelectionCost() throws Exception {
     }
     CountingList countedChildren = new CountingList(children);
     CollectOperator collect = new CollectOperator(new TestContext(), countedChildren);
+    collect.enableReadyInputConsumption();
     try {
       for (int round = 0; round < rounds; round++) {
         assertTrue(collect.isBlocked().isDone());
@@ -178,6 +183,7 @@ public void testClosedChildrenDoNotIncreaseRemainingChildSelectionCost() throws
     }
     CountingList countedChildren = new CountingList(children);
     CollectOperator collect = new CollectOperator(new TestContext(), countedChildren);
+    collect.enableReadyInputConsumption();
     try {
       // Retire the first 999 children and consume one block from the remaining child.
       for (int i = 0; i < childCount; i++) {
@@ -378,6 +384,47 @@ public void testCloseWakesWaitAndClosesEveryChildDespiteException() throws Excep
     assertEquals(1, b.closes);
   }
 
+  @Test
+  public void testDefaultCollectRetainsSequentialConsumption() throws Exception {
+    Child a = new Child(block(1));
+    Child b = new Child(block(2));
+    b.ready();
+    CollectOperator collect =
+        new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(a, b)));
+    assertFalse(collect.isBlocked().isDone());
+    a.ready();
+    assertSame(a.output, collect.next());
+    a.hasNext = false;
+    a.finished = true;
+    assertNull(collect.next());
+    assertSame(b.output, collect.next());
+    assertEquals(0, collect.calculateRetainedSizeAfterCallingNext());
+    collect.close();
+    assertEquals(1, a.closes);
+    assertEquals(1, b.closes);
+  }
+
+  @Test
+  public void testDefaultMappingCollectRetainsSequentialConsumption() throws Exception {
+    Child a = new Child(block(10, 11));
+    Child b = new Child(block(20, 21));
+    b.ready();
+    List<List<Integer>> mappings = new ArrayList<>();
+    mappings.add(Collections.singletonList(0));
+    mappings.add(Collections.singletonList(1));
+    MappingCollectOperator collect =
+        new MappingCollectOperator(
+            new TestContext(), new ArrayList<>(Arrays.asList(a, b)), mappings);
+    assertFalse(collect.isBlocked().isDone());
+    a.ready();
+    assertEquals(10, collect.next().getColumn(0).getInt(0));
+    a.hasNext = false;
+    a.finished = true;
+    assertNull(collect.next());
+    assertEquals(21, collect.next().getColumn(0).getInt(0));
+    collect.close();
+  }
+
   @Test
   public void testMappingUsesActualChildIdentity() throws Exception {
     Child a = new Child(block(10, 11));
@@ -389,6 +436,7 @@ public void testMappingUsesActualChildIdentity() throws Exception {
     MappingCollectOperator collect =
         new MappingCollectOperator(
             new TestContext(), new ArrayList<>(Arrays.asList(a, b)), mappings);
+    collect.enableReadyInputConsumption();
     assertTrue(collect.isBlocked().isDone());
     assertEquals(21, collect.next().getColumn(0).getInt(0));
     a.ready();
@@ -473,6 +521,7 @@ public void testConcurrentCompletionsConsumeEveryChildExactlyOnce() throws Excep
     List<Child> children = new ArrayList<>();
     for (int i = 0; i < count; i++) children.add(new Child(block(i)));
     CollectOperator collect = new CollectOperator(new TestContext(), new ArrayList<>(children));
+    collect.enableReadyInputConsumption();
     ListenableFuture<?> pending = collect.isBlocked();
     assertFalse(pending.isDone());
     ExecutorService executor = Executors.newFixedThreadPool(count);

iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java

@@ -3942,4 +3942,5 @@ private DataNodeQueryMessages() {}
 
   public static final String EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625 = "Plan node %s is already registered on memory side %s, not %s";
   public static final String EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F = "Memory reservation is not registered for query %s, fragment %s, plan node %s";
+  public static final String LOG_SHARED_EXCHANGE_BUDGET_FOR_FI_ARG_INPUT_ARG_OUTPUT_ARG_98AE1109 = "Shared exchange budget for FI {}: INPUT={}, OUTPUT={}";
 }

iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java

@@ -4698,4 +4698,5 @@ private DataNodeQueryMessages() {}
 
   public static final String EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625 = "PlanNode %s 已注册到内存侧 %s，不能注册到 %s";
   public static final String EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F = "查询 %s、Fragment %s、PlanNode %s 尚未注册内存账目";
+  public static final String LOG_SHARED_EXCHANGE_BUDGET_FOR_FI_ARG_INPUT_ARG_OUTPUT_ARG_98AE1109 = "FI {} 的共享 exchange 预算：INPUT={}，OUTPUT={}";
 }

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

@@ -22,6 +22,7 @@
 import org.apache.iotdb.calc.execution.filter.TopKRuntimeFilter;
 import org.apache.iotdb.calc.execution.operator.CommonOperatorContext;
 import org.apache.iotdb.calc.execution.operator.Operator;
+import org.apache.iotdb.calc.execution.operator.process.CollectOperator;
 import org.apache.iotdb.calc.execution.operator.process.FilterAndProjectOperator;
 import org.apache.iotdb.calc.execution.operator.process.LimitOperator;
 import org.apache.iotdb.calc.execution.operator.process.OffsetOperator;
@@ -47,7 +48,9 @@
 import org.apache.iotdb.commons.queryengine.plan.relational.metadata.QualifiedObjectName;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.Symbol;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.AggregationNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Expression;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.FunctionCall;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.LongLiteral;
@@ -301,6 +304,20 @@ protected OperatorContext addOperatorContext(
         .addOperatorContext(context.getNextOperatorId(), planNodeId, operatorType);
   }
 
+  @Override
+  public Operator visitCollect(CollectNode node, LocalExecutionPlanContext context) {
+    CollectOperator collect = (CollectOperator) super.visitCollect(node, context);
+    context.addCollectOperator(collect);
+    return collect;
+  }
+
+  @Override
+  public Operator visitUnion(UnionNode node, LocalExecutionPlanContext context) {
+    CollectOperator collect = (CollectOperator) super.visitUnion(node, context);
+    context.addCollectOperator(collect);
+    return collect;
+  }
+
   @Override
   public Operator visitCteScan(CteScanNode node, LocalExecutionPlanContext context) {
     OperatorContext operatorContext =