feat: Integrate MemoryReservationManager into percentile accumulators for improved memory management

Author: FearfulTomcat27

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/Percentile.java

@@ -45,7 +45,9 @@ public void addValue(double value) {
   }
 
   public void addValues(double... vals) {
-    if (vals == null || vals.length == 0) return;
+    if (vals == null || vals.length == 0) {
+      return;
+    }
 
     int newSize = size + vals.length;
     if (newSize > capacity) {

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/AccumulatorFactory.java

@@ -57,6 +57,7 @@
 import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.grouped.IntGroupedApproxMostFrequentAccumulator;
 import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.grouped.LongGroupedApproxMostFrequentAccumulator;
 import org.apache.iotdb.calc.i18n.CalcMessages;
+import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager;
 import org.apache.iotdb.common.rpc.thrift.TAggregationType;
 import org.apache.iotdb.commons.queryengine.execution.operator.source.relational.aggregation.grouped.UpdateMemory;
 import org.apache.iotdb.commons.queryengine.execution.operator.source.relational.aggregation.grouped.hash.MarkDistinctHash;
@@ -107,7 +108,8 @@ public static TableAccumulator createAccumulator(
       boolean isAggTableScan,
       String timeColumnName,
       Set<String> measurementColumnNames,
-      boolean distinct) {
+      boolean distinct,
+      MemoryReservationManager memoryReservationManager) {
     TableAccumulator result;
 
     // Input expression size of 1 indicates aggregation split has occurred and this is a final
@@ -167,7 +169,7 @@ public static TableAccumulator createAccumulator(
               ? new FirstAccumulator(inputDataTypes.get(0), isAggTableScan)
               : new FirstDescAccumulator(inputDataTypes.get(0));
     } else {
-      result = createBuiltinAccumulator(aggregationType, inputDataTypes);
+      result = createBuiltinAccumulator(aggregationType, inputDataTypes, memoryReservationManager);
     }
 
     if (distinct) {
@@ -189,7 +191,8 @@ public static GroupedAccumulator createGroupedAccumulator(
       List<Expression> inputExpressions,
       Map<String, String> inputAttributes,
       boolean ascending,
-      boolean distinct) {
+      boolean distinct,
+      MemoryReservationManager memoryReservationManager) {
     GroupedAccumulator result;
 
     if (aggregationType == TAggregationType.UDAF) {
@@ -198,7 +201,12 @@ public static GroupedAccumulator createGroupedAccumulator(
     } else {
       result =
           createBuiltinGroupedAccumulator(
-              aggregationType, inputDataTypes, inputExpressions, inputAttributes, ascending);
+              aggregationType,
+              inputDataTypes,
+              inputExpressions,
+              inputAttributes,
+              ascending,
+              memoryReservationManager);
     }
 
     if (distinct) {
@@ -243,7 +251,8 @@ private static GroupedAccumulator createBuiltinGroupedAccumulator(
       List<TSDataType> inputDataTypes,
       List<Expression> inputExpressions,
       Map<String, String> inputAttributes,
-      boolean ascending) {
+      boolean ascending,
+      MemoryReservationManager memoryReservationManager) {
     switch (aggregationType) {
       case COUNT:
         return new GroupedCountAccumulator();
@@ -328,15 +337,17 @@ private static GroupedAccumulator createBuiltinGroupedAccumulator(
         return new GroupedCentralMomentAccumulator(
             inputDataTypes.get(0), CentralMomentAccumulator.MomentType.KURTOSIS);
       case PERCENTILE:
-        return new GroupedPercentileAccumulator(inputDataTypes.get(0));
+        return new GroupedPercentileAccumulator(inputDataTypes.get(0), memoryReservationManager);
       default:
         throw new IllegalArgumentException(
             CalcMessages.INVALID_AGGREGATION_FUNCTION + aggregationType);
     }
   }
 
   public static TableAccumulator createBuiltinAccumulator(
-      TAggregationType aggregationType, List<TSDataType> inputDataTypes) {
+      TAggregationType aggregationType,
+      List<TSDataType> inputDataTypes,
+      MemoryReservationManager memoryReservationManager) {
     switch (aggregationType) {
       case COUNT:
         return new CountAccumulator();
@@ -422,7 +433,7 @@ public static TableAccumulator createBuiltinAccumulator(
         return new TableCentralMomentAccumulator(
             inputDataTypes.get(0), CentralMomentAccumulator.MomentType.KURTOSIS);
       case PERCENTILE:
-        return new PercentileAccumulator(inputDataTypes.get(0));
+        return new PercentileAccumulator(inputDataTypes.get(0), memoryReservationManager);
       default:
         throw new IllegalArgumentException(
             CalcMessages.INVALID_AGGREGATION_FUNCTION + aggregationType);

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/PercentileAccumulator.java

@@ -15,6 +15,7 @@
 package org.apache.iotdb.calc.execution.operator.source.relational.aggregation;
 
 import org.apache.iotdb.calc.execution.operator.source.relational.Percentile;
+import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager;
 import org.apache.iotdb.commons.exception.SemanticException;
 
 import org.apache.tsfile.block.column.Column;
@@ -36,8 +37,14 @@ public class PercentileAccumulator implements TableAccumulator {
   private Percentile percentile = new Percentile();
   private double percentage;
 
-  public PercentileAccumulator(TSDataType seriesDataType) {
+  private final MemoryReservationManager memoryReservationManager;
+  private long previousPercentileSize;
+
+  public PercentileAccumulator(
+      TSDataType seriesDataType, MemoryReservationManager memoryReservationManager) {
     this.seriesDataType = seriesDataType;
+    this.memoryReservationManager = memoryReservationManager;
+    updateMemoryReservation();
   }
 
   @Override
@@ -47,7 +54,7 @@ public long getEstimatedSize() {
 
   @Override
   public TableAccumulator copy() {
-    return new PercentileAccumulator(seriesDataType);
+    return new PercentileAccumulator(seriesDataType, memoryReservationManager);
   }
 
   @Override
@@ -60,21 +67,22 @@ public void addInput(Column[] arguments, AggregationMask mask) {
     switch (seriesDataType) {
       case INT32:
         addIntInput(arguments[0], mask);
-        return;
+        break;
       case INT64:
       case TIMESTAMP:
         addLongInput(arguments[0], mask);
-        return;
+        break;
       case FLOAT:
         addFloatInput(arguments[0], mask);
-        return;
+        break;
       case DOUBLE:
         addDoubleInput(arguments[0], mask);
-        return;
+        break;
       default:
         throw new UnSupportedDataTypeException(
             String.format("Unsupported data type in Percentile Aggregation: %s", seriesDataType));
     }
+    updateMemoryReservation();
   }
 
   @Override
@@ -87,6 +95,7 @@ public void addIntermediate(Column argument) {
         percentile.merge(Percentile.deserialize(buffer));
       }
     }
+    updateMemoryReservation();
   }
 
   @Override
@@ -137,7 +146,19 @@ public void addStatistics(Statistics[] statistics) {
 
   @Override
   public void reset() {
-    percentile.clear();
+    percentile = new Percentile();
+    updateMemoryReservation();
+  }
+
+  private void updateMemoryReservation() {
+    long currentSize = percentile.getEstimatedSize();
+    long delta = currentSize - previousPercentileSize;
+    if (delta > 0) {
+      memoryReservationManager.reserveMemoryCumulatively(delta);
+    } else if (delta < 0) {
+      memoryReservationManager.releaseMemoryCumulatively(-delta);
+    }
+    previousPercentileSize = currentSize;
   }
 
   private void addIntInput(Column column, AggregationMask mask) {

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/GroupedPercentileAccumulator.java

@@ -17,6 +17,7 @@
 import org.apache.iotdb.calc.execution.operator.source.relational.Percentile;
 import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.AggregationMask;
 import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.grouped.array.PercentileBigArray;
+import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager;
 import org.apache.iotdb.commons.exception.SemanticException;
 
 import org.apache.tsfile.block.column.Column;
@@ -34,10 +35,15 @@ public class GroupedPercentileAccumulator implements GroupedAccumulator {
       RamUsageEstimator.shallowSizeOfInstance(GroupedPercentileAccumulator.class);
   private final TSDataType seriesDataType;
   private double percentage;
+  private final MemoryReservationManager memoryReservationManager;
+  private long previousArraySize;
   private final PercentileBigArray array = new PercentileBigArray();
 
-  public GroupedPercentileAccumulator(TSDataType seriesDataType) {
+  public GroupedPercentileAccumulator(
+      TSDataType seriesDataType, MemoryReservationManager memoryReservationManager) {
     this.seriesDataType = seriesDataType;
+    this.memoryReservationManager = memoryReservationManager;
+    updateMemoryReservation();
   }
 
   @Override
@@ -76,6 +82,7 @@ public void addInput(int[] groupIds, Column[] arguments, AggregationMask mask) {
         throw new UnSupportedDataTypeException(
             String.format("Unsupported data type in PERCENTILE Aggregation: %s", seriesDataType));
     }
+    updateMemoryReservation();
   }
 
   @Override
@@ -90,6 +97,7 @@ public void addIntermediate(int[] groupIds, Column argument) {
         array.get(groupId).merge(other);
       }
     }
+    updateMemoryReservation();
   }
 
   @Override
@@ -136,6 +144,18 @@ public void prepareFinal() {}
   @Override
   public void reset() {
     array.reset();
+    updateMemoryReservation();
+  }
+
+  private void updateMemoryReservation() {
+    long currentSize = array.sizeOf();
+    long delta = currentSize - previousArraySize;
+    if (delta > 0) {
+      memoryReservationManager.reserveMemoryCumulatively(delta);
+    } else if (delta < 0) {
+      memoryReservationManager.releaseMemoryCumulatively(-delta);
+    }
+    previousArraySize = currentSize;
   }
 
   public void addIntInput(int[] groupIds, Column[] arguments, AggregationMask mask) {