Adapt ready-first collect to Apache master

Author: ColinLeeo

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

@@ -91,7 +91,9 @@ public ListenableFuture<?> isBlocked() {
     if (!inited) {
       inited = true;
       for (Operator child : children) {
-        child.isBlocked();
+        if (child != null) {
+          child.isBlocked();
+        }
       }
     }
     if (children.isEmpty()) {

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/plan/planner/memory/MemoryReservationManager.java

@@ -54,6 +54,11 @@ public interface MemoryReservationManager {
    */
   void releaseMemoryImmediately(final long size);
 
+  /** Flushes any batched release accumulated by {@link #releaseMemoryCumulatively(long)}. */
+  default void releaseMemoryImmediately() {
+    // Implementations that batch releases override this method.
+  }
+
   /**
    * Release all reserved memory immediately. Make sure this method is called when the lifecycle of
    * this manager ends, Or the memory to be released in the batch may not be released correctly.

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java

@@ -159,7 +159,7 @@ public void setMemoryReservationGroupId(String reservationPlanNodeId) {
       // Collect group. Restore the fragment-instance limit so the shared reservation key, rather
       // than an individual queue, is the group-wide limit.
       this.maxBytesCanReserve =
-          IoTDBDescriptor.getInstance().getConfig().getMaxBytesPerFragmentInstance();
+          IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance();
     }
   }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java

@@ -567,7 +567,7 @@ public synchronized void setMemoryReservationGroupId(String reservationPlanNodeI
       // Collect group. Restore the fragment-instance limit so the shared reservation key, rather
       // than an individual handle, is the group-wide limit.
       this.maxBytesCanReserve =
-          IoTDBDescriptor.getInstance().getConfig().getMaxBytesPerFragmentInstance();
+          IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance();
     }
     localMemoryManager
         .getQueryPool()

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

@@ -47,6 +47,7 @@
 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.sql.ast.Expression;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.FunctionCall;
@@ -229,6 +230,16 @@ public DataNodeTableOperatorGenerator(Metadata metadata) {
     super(metadata);
   }
 
+  @Override
+  public Operator visitCollect(CollectNode node, LocalExecutionPlanContext context) {
+    CollectQuotaGroup quotaGroup = context.enterCollectQuotaGroup(node.getPlanNodeId().getId());
+    try {
+      return super.visitCollect(node, context);
+    } finally {
+      context.exitCollectQuotaGroup(quotaGroup);
+    }
+  }
+
   /**
    * Producer side: register a shared {@link TopKRuntimeFilter} before visiting TopK children so
    * nested Scan nodes can resolve the same instance during operator generation.
@@ -385,6 +396,7 @@ public Operator visitTableExchange(ExchangeNode node, LocalExecutionPlanContext
     if (!isSameNode) {
       context.addExchangeSumNum(1);
     }
+    context.registerExchangeSourceHandle(sourceHandle);
     sourceHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve());
     ExchangeOperator exchangeOperator =
         new ExchangeOperator(operatorContext, sourceHandle, node.getUpstreamPlanNodeId());

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

@@ -299,7 +299,8 @@ public void setMaxBytesOneHandleCanReserve() {
         exchangeSumNum);
     java.util.Set<ISourceHandle> groupedSourceHandles =
         Collections.newSetFromMap(new IdentityHashMap<>());
-    long maxBytesPerFI = IoTDBDescriptor.getInstance().getConfig().getMaxBytesPerFragmentInstance();
+    long maxBytesPerFI =
+        IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance();
     for (CollectQuotaGroup group : collectQuotaGroups) {
       for (ISourceHandle sourceHandle : group.getSourceHandles()) {
         groupedSourceHandles.add(sourceHandle);

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

@@ -1181,7 +1181,13 @@ public Operator visitCollect(CollectNode node, LocalExecutionPlanContext context
                 context.getNextOperatorId(),
                 node.getPlanNodeId(),
                 CollectOperator.class.getSimpleName());
-    return new CollectOperator(operatorContext, dealWithConsumeChildrenOneByOneNode(node, context));
+    CollectQuotaGroup quotaGroup = context.enterCollectQuotaGroup(node.getPlanNodeId().getId());
+    try {
+      return new CollectOperator(
+          operatorContext, dealWithConsumeChildrenOneByOneNode(node, context));
+    } finally {
+      context.exitCollectQuotaGroup(quotaGroup);
+    }
   }
 
   @Override

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

@@ -143,6 +143,14 @@ public void releaseMemoryImmediately(final long size) {
     }
   }
 
+  @Override
+  public void releaseMemoryImmediately() {
+    if (bytesToBeReleased > 0) {
+      releaseBytesImmediately(bytesToBeReleased);
+      bytesToBeReleased = 0;
+    }
+  }
+
   private void releaseBytesImmediately(final long size) {
     long poolBytes = deductReleaseAccounting(size);
     if (poolBytes > 0) {

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

@@ -56,6 +56,11 @@ public synchronized void releaseMemoryImmediately(long size) {
     super.releaseMemoryImmediately(size);
   }
 
+  @Override
+  public synchronized void releaseMemoryImmediately() {
+    super.releaseMemoryImmediately();
+  }
+
   @Override
   public synchronized void releaseAllReservedMemory() {
     super.releaseAllReservedMemory();

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/CollectOperatorTest.java

@@ -19,7 +19,9 @@
 
 package org.apache.iotdb.db.queryengine.execution.operator;
 
-import org.apache.iotdb.db.queryengine.execution.operator.process.CollectOperator;
+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 com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.SettableFuture;
@@ -81,7 +83,7 @@ private static TsBlock block(int value) {
   private static final class TestOperator implements Operator {
     private final List<TsBlock> blocks;
     private final ListenableFuture<?> blocked;
-    private final OperatorContext operatorContext = Mockito.mock(OperatorContext.class);
+    private final CommonOperatorContext operatorContext = Mockito.mock(CommonOperatorContext.class);
     private int index;
 
     private TestOperator(TsBlock block, ListenableFuture<?> blocked) {
@@ -94,7 +96,7 @@ private TestOperator(List<TsBlock> blocks, ListenableFuture<?> blocked) {
     }
 
     @Override
-    public OperatorContext getOperatorContext() {
+    public CommonOperatorContext getOperatorContext() {
       return operatorContext;
     }