diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/AggregationDistributionTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/AggregationDistributionTest.java index 75467533b8ea..3618f99fafa0 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/AggregationDistributionTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/AggregationDistributionTest.java @@ -362,29 +362,35 @@ public void testAggregationWithMultiGroupByLevelNode() throws IllegalPathExcepti fragmentInstances.forEach( f -> verifyAggregationStep(expectedStep, f.getFragment().getPlanNodeTree())); - Map> expectedDescriptorValue = new HashMap<>(); - expectedDescriptorValue.put(groupedPath, Collections.singletonList(groupedPath)); - assertTrue( - fragmentInstances - .get(0) - .getFragment() - .getPlanNodeTree() - .getChildren() - .get(0) - .getChildren() - .get(0) - instanceof GroupByLevelNode); - verifyGroupByLevelDescriptor( - expectedDescriptorValue, - (GroupByLevelNode) - fragmentInstances.get(0).getFragment().getPlanNodeTree().getChildren().get(0)); + Map> expectedDescriptorValue1 = new HashMap<>(); + expectedDescriptorValue1.put(groupedPath, Collections.singletonList(groupedPath)); Map> expectedDescriptorValue2 = new HashMap<>(); expectedDescriptorValue2.put(groupedPath, Arrays.asList(d3s1Path, d4s1Path)); - verifyGroupByLevelDescriptor( - expectedDescriptorValue2, - (GroupByLevelNode) - fragmentInstances.get(1).getFragment().getPlanNodeTree().getChildren().get(0)); + + boolean foundFirst = false; + boolean foundSecond = false; + + for (FragmentInstance instance : fragmentInstances) { + PlanNode planNodeTree = instance.getFragment().getPlanNodeTree(); + if (planNodeTree.getChildren().isEmpty()) { + continue; + } + PlanNode childNode = planNodeTree.getChildren().get(0); + if (!(childNode instanceof GroupByLevelNode)) { + continue; + } + + GroupByLevelNode groupByLevel = (GroupByLevelNode) childNode; + if (matchesExpectedDescriptor(groupByLevel, expectedDescriptorValue1)) { + foundFirst = true; + } else if (matchesExpectedDescriptor(groupByLevel, expectedDescriptorValue2)) { + foundSecond = true; + } + } + + assertTrue("Expected to find fragment with single grouped path descriptor", foundFirst); + assertTrue("Expected to find fragment with two specific paths descriptor", foundSecond); } @Test @@ -857,6 +863,28 @@ private void verifyGroupByLevelDescriptor( } } + private boolean matchesExpectedDescriptor( + GroupByLevelNode node, Map> expected) { + List descriptors = node.getGroupByLevelDescriptors(); + if (descriptors.size() != expected.size()) { + return false; + } + for (CrossSeriesAggregationDescriptor descriptor : descriptors) { + String outputExpression = descriptor.getOutputExpressions().get(0).getExpressionString(); + List expectedInputs = expected.get(outputExpression); + if (expectedInputs == null + || expectedInputs.size() != descriptor.getInputExpressions().size()) { + return false; + } + for (Expression inputExpression : descriptor.getInputExpressions()) { + if (!expectedInputs.contains(inputExpression.getExpressionString())) { + return false; + } + } + } + return true; + } + private void verifySlidingWindowDescriptor( List expected, SlidingWindowAggregationNode node) { List descriptorList = node.getAggregationDescriptorList();