From f81ec38440ca61bde6bf358698deff268c64d8a2 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:36:31 +0800 Subject: [PATCH] Fix aggregation with finished child input --- .../execution/aggregation/TreeAggregator.java | 20 ++- .../operator/process/AggregationOperator.java | 3 + .../operator/AggregationOperatorTest.java | 145 ++++++++++++++++++ 3 files changed, 163 insertions(+), 5 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/aggregation/TreeAggregator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/aggregation/TreeAggregator.java index 41dee7116df61..43dcfb30ad7b1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/aggregation/TreeAggregator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/aggregation/TreeAggregator.java @@ -92,17 +92,27 @@ public void processTsBlocks(TsBlock[] tsBlock) { checkArgument( inputLocationList.size() == 1, DataNodeQueryMessages.EXCEPTION_FINAL_OUTPUT_CAN_ONLY_BE_SINGLE_COLUMN_6D82F9E0); + InputLocation inputLocation = inputLocationList.get(0)[0]; + if (tsBlock[inputLocation.getTsBlockIndex()] == null) { + return; + } Column finalResult = - tsBlock[inputLocationList.get(0)[0].getTsBlockIndex()].getColumn( - inputLocationList.get(0)[0].getValueColumnIndex()); + tsBlock[inputLocation.getTsBlockIndex()].getColumn(inputLocation.getValueColumnIndex()); accumulator.setFinal(finalResult); } else { for (InputLocation[] inputLocations : inputLocationList) { Column[] columns = new Column[inputLocations.length]; + boolean hasMissingInput = false; for (int i = 0; i < inputLocations.length; i++) { - columns[i] = - tsBlock[inputLocations[i].getTsBlockIndex()].getColumn( - inputLocations[i].getValueColumnIndex()); + TsBlock inputTsBlock = tsBlock[inputLocations[i].getTsBlockIndex()]; + if (inputTsBlock == null) { + hasMissingInput = true; + break; + } + columns[i] = inputTsBlock.getColumn(inputLocations[i].getValueColumnIndex()); + } + if (hasMissingInput) { + continue; } accumulator.addIntermediate(columns); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/AggregationOperator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/AggregationOperator.java index eec1c01743092..0b954046e19fd 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/AggregationOperator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/AggregationOperator.java @@ -161,6 +161,9 @@ private void calculateNextAggregationResult() { } for (int i = 0; i < inputOperatorsCount; i++) { + if (inputTsBlocks[i] == null) { + continue; + } inputTsBlocks[i] = inputTsBlocks[i].skipFirst(); if (inputTsBlocks[i].isEmpty()) { inputTsBlocks[i] = null; diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/AggregationOperatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/AggregationOperatorTest.java index a32d0883b4ebb..fc79146502ca1 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/AggregationOperatorTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/AggregationOperatorTest.java @@ -53,6 +53,7 @@ import org.apache.tsfile.exception.write.WriteProcessException; import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.read.common.block.TsBlock; +import org.apache.tsfile.read.common.block.TsBlockBuilder; import org.apache.tsfile.utils.TimeDuration; import org.apache.tsfile.write.schema.IMeasurementSchema; import org.apache.tsfile.write.schema.MeasurementSchema; @@ -292,6 +293,85 @@ public void testGroupByIntermediateResult2() throws Exception { assertEquals(4, count); } + @Test + public void testGroupByIntermediateResultWithFinishedChild() throws Exception { + QueryId queryId = new QueryId("stub_query_finished_child"); + FragmentInstanceId instanceId = + new FragmentInstanceId(new PlanFragmentId(queryId, 0), "stub-instance"); + FragmentInstanceStateMachine stateMachine = + new FragmentInstanceStateMachine(instanceId, instanceNotificationExecutor); + FragmentInstanceContext fragmentInstanceContext = + createFragmentInstanceContext(instanceId, stateMachine); + DriverContext driverContext = new DriverContext(fragmentInstanceContext, 0); + driverContext.addOperatorContext( + 1, new PlanNodeId("finished-child-1"), FixedTsBlockOperator.class.getSimpleName()); + driverContext.addOperatorContext( + 2, new PlanNodeId("finished-child-2"), FixedTsBlockOperator.class.getSimpleName()); + driverContext.addOperatorContext( + 3, new PlanNodeId("aggregation"), AggregationOperator.class.getSimpleName()); + driverContext + .getOperatorContexts() + .forEach(operatorContext -> OperatorContext.setMaxRunTime(TEST_TIME_SLICE)); + + List children = new ArrayList<>(); + children.add( + new FixedTsBlockOperator( + driverContext.getOperatorContexts().get(0), buildCountTsBlock(new long[] {0}, 3))); + children.add( + new FixedTsBlockOperator( + driverContext.getOperatorContexts().get(1), + buildCountTsBlock(new long[] {0, 100}, 5, 7))); + + List inputLocationForCount = new ArrayList<>(); + inputLocationForCount.add(new InputLocation[] {new InputLocation(0, 0)}); + inputLocationForCount.add(new InputLocation[] {new InputLocation(1, 0)}); + + List finalAggregators = new ArrayList<>(); + finalAggregators.add( + new TreeAggregator( + AccumulatorFactory.createBuiltinAccumulators( + Collections.singletonList(TAggregationType.COUNT), + TSDataType.INT32, + Collections.emptyList(), + Collections.emptyMap(), + true) + .get(0), + AggregationStep.FINAL, + inputLocationForCount)); + + GroupByTimeParameter groupByTimeParameter = + new GroupByTimeParameter(0, 200, new TimeDuration(0, 100), new TimeDuration(0, 100), true); + AggregationOperator aggregationOperator = + new AggregationOperator( + driverContext.getOperatorContexts().get(2), + finalAggregators, + initTimeRangeIterator(groupByTimeParameter, true, true, ZoneId.systemDefault()), + children, + false, + DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES); + + long[] expectedTimes = new long[] {0, 100}; + long[] expectedCounts = new long[] {8, 7}; + int count = 0; + while (true) { + ListenableFuture blocked = aggregationOperator.isBlocked(); + blocked.get(); + if (!aggregationOperator.hasNext()) { + break; + } + TsBlock resultTsBlock = aggregationOperator.next(); + if (resultTsBlock == null) { + continue; + } + for (int pos = 0; pos < resultTsBlock.getPositionCount(); pos++) { + assertEquals(expectedTimes[count], resultTsBlock.getTimeColumn().getLong(pos)); + assertEquals(expectedCounts[count], resultTsBlock.getColumn(0).getLong(pos)); + count++; + } + } + assertEquals(2, count); + } + /** * @param aggregationTypes Aggregation function used in test * @param groupByTimeParameter group by time parameter @@ -410,4 +490,69 @@ private AggregationOperator initAggregationOperator( false, DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES); } + + private TsBlock buildCountTsBlock(long[] times, long... counts) { + TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT64)); + for (int i = 0; i < times.length; i++) { + builder.getTimeColumnBuilder().writeLong(times[i]); + builder.getColumnBuilder(0).writeLong(counts[i]); + builder.declarePosition(); + } + return builder.build(); + } + + private static class FixedTsBlockOperator implements Operator { + + private final OperatorContext operatorContext; + private final TsBlock[] tsBlocks; + private int index; + + private FixedTsBlockOperator(OperatorContext operatorContext, TsBlock... tsBlocks) { + this.operatorContext = operatorContext; + this.tsBlocks = tsBlocks; + } + + @Override + public OperatorContext getOperatorContext() { + return operatorContext; + } + + @Override + public TsBlock next() { + return tsBlocks[index++]; + } + + @Override + public boolean hasNext() { + return index < tsBlocks.length; + } + + @Override + public void close() {} + + @Override + public boolean isFinished() { + return !hasNext(); + } + + @Override + public long calculateMaxPeekMemory() { + return 0; + } + + @Override + public long calculateMaxReturnSize() { + return 0; + } + + @Override + public long calculateRetainedSizeAfterCallingNext() { + return 0; + } + + @Override + public long ramBytesUsed() { + return 0; + } + } }