Skip to content

Commit 87cc668

Browse files
FearfulTomcat27FearfulTomcat27
committed
feat: Integrate MemoryReservationManager into percentile accumulators for improved memory management
1 parent 7f55be6 commit 87cc668

9 files changed

Lines changed: 143 additions & 39 deletions

File tree

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,9 @@ public void addValue(double value) {
4545
}
4646

4747
public void addValues(double... vals) {
48-
if (vals == null || vals.length == 0) return;
48+
if (vals == null || vals.length == 0) {
49+
return;
50+
}
4951

5052
int newSize = size + vals.length;
5153
if (newSize > capacity) {

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

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.grouped.IntGroupedApproxMostFrequentAccumulator;
5151
import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.grouped.LongGroupedApproxMostFrequentAccumulator;
5252
import org.apache.iotdb.calc.i18n.CalcMessages;
53+
import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager;
5354
import org.apache.iotdb.common.rpc.thrift.TAggregationType;
5455
import org.apache.iotdb.commons.queryengine.execution.operator.source.relational.aggregation.grouped.UpdateMemory;
5556
import org.apache.iotdb.commons.queryengine.execution.operator.source.relational.aggregation.grouped.hash.MarkDistinctHash;
@@ -100,7 +101,8 @@ public static TableAccumulator createAccumulator(
100101
boolean isAggTableScan,
101102
String timeColumnName,
102103
Set<String> measurementColumnNames,
103-
boolean distinct) {
104+
boolean distinct,
105+
MemoryReservationManager memoryReservationManager) {
104106
TableAccumulator result;
105107

106108
// Input expression size of 1 indicates aggregation split has occurred and this is a final
@@ -160,7 +162,7 @@ public static TableAccumulator createAccumulator(
160162
? new FirstAccumulator(inputDataTypes.get(0), isAggTableScan)
161163
: new FirstDescAccumulator(inputDataTypes.get(0));
162164
} else {
163-
result = createBuiltinAccumulator(aggregationType, inputDataTypes);
165+
result = createBuiltinAccumulator(aggregationType, inputDataTypes, memoryReservationManager);
164166
}
165167

166168
if (distinct) {
@@ -182,7 +184,8 @@ public static GroupedAccumulator createGroupedAccumulator(
182184
List<Expression> inputExpressions,
183185
Map<String, String> inputAttributes,
184186
boolean ascending,
185-
boolean distinct) {
187+
boolean distinct,
188+
MemoryReservationManager memoryReservationManager) {
186189
GroupedAccumulator result;
187190

188191
if (aggregationType == TAggregationType.UDAF) {
@@ -191,7 +194,12 @@ public static GroupedAccumulator createGroupedAccumulator(
191194
} else {
192195
result =
193196
createBuiltinGroupedAccumulator(
194-
aggregationType, inputDataTypes, inputExpressions, inputAttributes, ascending);
197+
aggregationType,
198+
inputDataTypes,
199+
inputExpressions,
200+
inputAttributes,
201+
ascending,
202+
memoryReservationManager);
195203
}
196204

197205
if (distinct) {
@@ -236,7 +244,8 @@ private static GroupedAccumulator createBuiltinGroupedAccumulator(
236244
List<TSDataType> inputDataTypes,
237245
List<Expression> inputExpressions,
238246
Map<String, String> inputAttributes,
239-
boolean ascending) {
247+
boolean ascending,
248+
MemoryReservationManager memoryReservationManager) {
240249
switch (aggregationType) {
241250
case COUNT:
242251
return new GroupedCountAccumulator();
@@ -293,15 +302,17 @@ private static GroupedAccumulator createBuiltinGroupedAccumulator(
293302
return new GroupedApproxPercentileWithWeightAccumulator(inputDataTypes.get(0));
294303
}
295304
case PERCENTILE:
296-
return new GroupedPercentileAccumulator(inputDataTypes.get(0));
305+
return new GroupedPercentileAccumulator(inputDataTypes.get(0), memoryReservationManager);
297306
default:
298307
throw new IllegalArgumentException(
299308
CalcMessages.INVALID_AGGREGATION_FUNCTION + aggregationType);
300309
}
301310
}
302311

303312
public static TableAccumulator createBuiltinAccumulator(
304-
TAggregationType aggregationType, List<TSDataType> inputDataTypes) {
313+
TAggregationType aggregationType,
314+
List<TSDataType> inputDataTypes,
315+
MemoryReservationManager memoryReservationManager) {
305316
switch (aggregationType) {
306317
case COUNT:
307318
return new CountAccumulator();
@@ -359,7 +370,7 @@ public static TableAccumulator createBuiltinAccumulator(
359370
return new ApproxPercentileWithWeightAccumulator(inputDataTypes.get(0));
360371
}
361372
case PERCENTILE:
362-
return new PercentileAccumulator(inputDataTypes.get(0));
373+
return new PercentileAccumulator(inputDataTypes.get(0), memoryReservationManager);
363374
default:
364375
throw new IllegalArgumentException(
365376
CalcMessages.INVALID_AGGREGATION_FUNCTION + aggregationType);

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

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
package org.apache.iotdb.calc.execution.operator.source.relational.aggregation;
1616

1717
import org.apache.iotdb.calc.execution.operator.source.relational.Percentile;
18+
import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager;
1819
import org.apache.iotdb.commons.exception.SemanticException;
1920

2021
import org.apache.tsfile.block.column.Column;
@@ -36,8 +37,14 @@ public class PercentileAccumulator implements TableAccumulator {
3637
private Percentile percentile = new Percentile();
3738
private double percentage;
3839

39-
public PercentileAccumulator(TSDataType seriesDataType) {
40+
private final MemoryReservationManager memoryReservationManager;
41+
private long previousPercentileSize;
42+
43+
public PercentileAccumulator(
44+
TSDataType seriesDataType, MemoryReservationManager memoryReservationManager) {
4045
this.seriesDataType = seriesDataType;
46+
this.memoryReservationManager = memoryReservationManager;
47+
updateMemoryReservation();
4148
}
4249

4350
@Override
@@ -47,7 +54,7 @@ public long getEstimatedSize() {
4754

4855
@Override
4956
public TableAccumulator copy() {
50-
return new PercentileAccumulator(seriesDataType);
57+
return new PercentileAccumulator(seriesDataType, memoryReservationManager);
5158
}
5259

5360
@Override
@@ -60,21 +67,22 @@ public void addInput(Column[] arguments, AggregationMask mask) {
6067
switch (seriesDataType) {
6168
case INT32:
6269
addIntInput(arguments[0], mask);
63-
return;
70+
break;
6471
case INT64:
6572
case TIMESTAMP:
6673
addLongInput(arguments[0], mask);
67-
return;
74+
break;
6875
case FLOAT:
6976
addFloatInput(arguments[0], mask);
70-
return;
77+
break;
7178
case DOUBLE:
7279
addDoubleInput(arguments[0], mask);
73-
return;
80+
break;
7481
default:
7582
throw new UnSupportedDataTypeException(
7683
String.format("Unsupported data type in Percentile Aggregation: %s", seriesDataType));
7784
}
85+
updateMemoryReservation();
7886
}
7987

8088
@Override
@@ -87,6 +95,7 @@ public void addIntermediate(Column argument) {
8795
percentile.merge(Percentile.deserialize(buffer));
8896
}
8997
}
98+
updateMemoryReservation();
9099
}
91100

92101
@Override
@@ -137,7 +146,19 @@ public void addStatistics(Statistics[] statistics) {
137146

138147
@Override
139148
public void reset() {
140-
percentile.clear();
149+
percentile = new Percentile();
150+
updateMemoryReservation();
151+
}
152+
153+
private void updateMemoryReservation() {
154+
long currentSize = percentile.getEstimatedSize();
155+
long delta = currentSize - previousPercentileSize;
156+
if (delta > 0) {
157+
memoryReservationManager.reserveMemoryCumulatively(delta);
158+
} else if (delta < 0) {
159+
memoryReservationManager.releaseMemoryCumulatively(-delta);
160+
}
161+
previousPercentileSize = currentSize;
141162
}
142163

143164
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‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import org.apache.iotdb.calc.execution.operator.source.relational.Percentile;
1818
import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.AggregationMask;
1919
import org.apache.iotdb.calc.execution.operator.source.relational.aggregation.grouped.array.PercentileBigArray;
20+
import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager;
2021
import org.apache.iotdb.commons.exception.SemanticException;
2122

2223
import org.apache.tsfile.block.column.Column;
@@ -34,10 +35,15 @@ public class GroupedPercentileAccumulator implements GroupedAccumulator {
3435
RamUsageEstimator.shallowSizeOfInstance(GroupedPercentileAccumulator.class);
3536
private final TSDataType seriesDataType;
3637
private double percentage;
38+
private final MemoryReservationManager memoryReservationManager;
39+
private long previousArraySize;
3740
private final PercentileBigArray array = new PercentileBigArray();
3841

39-
public GroupedPercentileAccumulator(TSDataType seriesDataType) {
42+
public GroupedPercentileAccumulator(
43+
TSDataType seriesDataType, MemoryReservationManager memoryReservationManager) {
4044
this.seriesDataType = seriesDataType;
45+
this.memoryReservationManager = memoryReservationManager;
46+
updateMemoryReservation();
4147
}
4248

4349
@Override
@@ -76,6 +82,7 @@ public void addInput(int[] groupIds, Column[] arguments, AggregationMask mask) {
7682
throw new UnSupportedDataTypeException(
7783
String.format("Unsupported data type in PERCENTILE Aggregation: %s", seriesDataType));
7884
}
85+
updateMemoryReservation();
7986
}
8087

8188
@Override
@@ -90,6 +97,7 @@ public void addIntermediate(int[] groupIds, Column argument) {
9097
array.get(groupId).merge(other);
9198
}
9299
}
100+
updateMemoryReservation();
93101
}
94102

95103
@Override
@@ -136,6 +144,18 @@ public void prepareFinal() {}
136144
@Override
137145
public void reset() {
138146
array.reset();
147+
updateMemoryReservation();
148+
}
149+
150+
private void updateMemoryReservation() {
151+
long currentSize = array.sizeOf();
152+
long delta = currentSize - previousArraySize;
153+
if (delta > 0) {
154+
memoryReservationManager.reserveMemoryCumulatively(delta);
155+
} else if (delta < 0) {
156+
memoryReservationManager.releaseMemoryCumulatively(-delta);
157+
}
158+
previousArraySize = currentSize;
139159
}
140160

141161
public void addIntInput(int[] groupIds, Column[] arguments, AggregationMask mask) {

0 commit comments

Comments
 (0)