Stabilize query ordering and template column pagination

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/visitor/SchemaTreeVisitor.java

@@ -29,16 +29,20 @@
 import org.apache.tsfile.write.schema.IMeasurementSchema;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.TreeMap;
 
 import static org.apache.iotdb.commons.schema.SchemaConstant.NON_TEMPLATE;
 
 public abstract class SchemaTreeVisitor<R> extends AbstractTreeVisitor<SchemaNode, R> {
 
   private Map<Integer, Template> templateMap;
+  // Reuse the ordered snapshot for devices that share a template.
+  private final Map<Integer, List<IMeasurementSchema>> templateMeasurements = new HashMap<>();
 
   protected SchemaTreeVisitor() {}
 
@@ -63,6 +67,7 @@ public List<R> getAllResult() {
 
   public void setTemplateMap(Map<Integer, Template> templateMap) {
     this.templateMap = templateMap;
+    templateMeasurements.clear();
   }
 
   @Override
@@ -123,10 +128,12 @@ protected Iterator<SchemaNode> getChildrenIterator(SchemaNode parent) {
         && parent.getAsEntityNode().getTemplateId() != NON_TEMPLATE) {
       return new Iterator<SchemaNode>() {
         private final Iterator<IMeasurementSchema> templateChildrenIterator =
-            templateMap
-                .get(parent.getAsEntityNode().getTemplateId())
-                .getSchemaMap()
-                .values()
+            templateMeasurements
+                .computeIfAbsent(
+                    parent.getAsEntityNode().getTemplateId(),
+                    templateId ->
+                        new ArrayList<>(
+                            new TreeMap<>(templateMap.get(templateId).getSchemaMap()).values()))
                 .iterator();
         private final Iterator<SchemaNode> directChildrenIterator = parent.getChildrenIterator();
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/schematree/visitor/SchemaTreeVisitorWithLimitOffsetWrapper.java

@@ -42,13 +42,6 @@ public SchemaTreeVisitorWithLimitOffsetWrapper(
     this.limit = limit;
     this.offset = offset;
     hasLimit = limit > 0 || offset > 0;
-
-    if (hasLimit) {
-      while (curOffset < offset && visitor.hasNext()) {
-        visitor.next();
-        curOffset++;
-      }
-    }
   }
 
   public void setTemplateMap(Map<Integer, Template> templateMap) {
@@ -57,11 +50,12 @@ public void setTemplateMap(Map<Integer, Template> templateMap) {
 
   @Override
   public boolean hasNext() {
-    if (hasLimit) {
-      return count < limit && visitor.hasNext();
-    } else {
-      return visitor.hasNext();
+    // Template metadata is set after construction, so defer traversal until iteration starts.
+    while (curOffset < offset && visitor.hasNext()) {
+      visitor.next();
+      curOffset++;
     }
+    return (limit == 0 || count < limit) && visitor.hasNext();
   }
 
   @Override
@@ -119,12 +113,6 @@ public void reset() {
     visitor.reset();
     count = 0;
     curOffset = 0;
-    if (hasLimit) {
-      while (curOffset < offset && visitor.hasNext()) {
-        visitor.next();
-        curOffset++;
-      }
-    }
   }
 
   public int getNextOffset() {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java

@@ -1214,7 +1214,7 @@ private void analyzeGroupByTag(
 
     List<String> tagKeys = queryStatement.getGroupByTagComponent().getTagKeys();
     Map<List<String>, LinkedHashMap<Expression, List<Expression>>>
-        tagValuesToGroupedTimeseriesOperands = new HashMap<>();
+        tagValuesToGroupedTimeseriesOperands = new LinkedHashMap<>();
     LinkedHashMap<Expression, Set<Expression>> outputExpressionToRawExpressionsMap =
         new LinkedHashMap<>();
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/TemplatedAggregationAnalyze.java

@@ -45,6 +45,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
 
 import static org.apache.iotdb.calc.utils.constant.SqlConstant.COUNT_TIME;
 import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.ENDTIME;
@@ -106,6 +107,11 @@ static boolean canBuildAggregationPlanUseTemplate(
       return false;
     }
 
+    if (outputExpressions.isEmpty()) {
+      analysis.setFinishQueryAfterAnalyze(true);
+      return true;
+    }
+
     analyzeDeviceToWhere(analysis, queryStatement);
     if (deviceList.isEmpty()) {
       analysis.setFinishQueryAfterAnalyze(true);
@@ -180,6 +186,13 @@ private static boolean analyzeSelect(
       if (selectExpression instanceof FunctionExpression
           && COUNT_TIME.equalsIgnoreCase(
               ((FunctionExpression) selectExpression).getFunctionName())) {
+        if (paginationController.hasCurOffset()) {
+          paginationController.consumeOffset();
+          continue;
+        }
+        if (!paginationController.hasCurLimit()) {
+          break;
+        }
         FunctionExpression countTimeExpression = (FunctionExpression) selectExpression;
         // The parsed count_time(*) expression belongs to the statement and may be analyzed again
         // after a dispatch failure. Build a template-only expression instead of replacing its
@@ -203,6 +216,7 @@ private static boolean analyzeSelect(
         analysis
             .getExpressionTypes()
             .put(NodeRef.of(templatedCountTimeExpression), TSDataType.INT64);
+        paginationController.consumeLimit();
         continue;
       }
 
@@ -214,7 +228,9 @@ private static boolean analyzeSelect(
             && "*".equalsIgnoreCase(selectExpression.getExpressions().get(0).getOutputSymbol())) {
           subExpressions = new ArrayList<>();
           FunctionExpression functionExpression = (FunctionExpression) selectExpression;
-          for (Map.Entry<String, IMeasurementSchema> entry : template.getSchemaMap().entrySet()) {
+          // Use the same stable wildcard order as raw-data queries, before applying pagination.
+          for (Map.Entry<String, IMeasurementSchema> entry :
+              new TreeMap<>(template.getSchemaMap()).entrySet()) {
             FunctionExpression subFunctionExpression =
                 new FunctionExpression(
                     functionExpression.getFunctionName(),
@@ -241,6 +257,7 @@ private static boolean analyzeSelect(
           selectExpressions.add(expression);
           aggregationExpressions.add(expression);
           analyzeExpressionType(analysis, expression);
+          paginationController.consumeLimit();
         } else {
           break;
         }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/TemplatedAnalyze.java

@@ -56,6 +56,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
 import java.util.stream.Collectors;
 
 import static org.apache.iotdb.db.queryengine.metric.QueryPlanCostMetricSet.PARTITION_FETCHER;
@@ -134,7 +135,9 @@ public static boolean canBuildPlanUseTemplate(
     for (ResultColumn resultColumn : queryStatement.getSelectComponent().getResultColumns()) {
       Expression expression = resultColumn.getExpression();
       if ("*".equals(expression.getOutputSymbol())) {
-        for (Map.Entry<String, IMeasurementSchema> entry : template.getSchemaMap().entrySet()) {
+        // Template map iteration can change after serialization. Order columns before pagination.
+        for (Map.Entry<String, IMeasurementSchema> entry :
+            new TreeMap<>(template.getSchemaMap()).entrySet()) {
           if (paginationController.hasCurOffset()) {
             paginationController.consumeOffset();
           } else if (paginationController.hasCurLimit()) {
@@ -168,6 +171,7 @@ public static boolean canBuildPlanUseTemplate(
             // reserve memory for this expression
             context.reserveMemoryForFrontEnd(measurementPath.ramBytesUsed());
             outputExpressions.add(new Pair<>(measurementPath, resultColumn.getAlias()));
+            paginationController.consumeLimit();
           } else {
             break;
           }
@@ -183,6 +187,11 @@ public static boolean canBuildPlanUseTemplate(
 
     analyzeSelect(queryStatement, analysis, outputExpressions, template);
 
+    if (outputExpressions.isEmpty()) {
+      analysis.setFinishQueryAfterAnalyze(true);
+      return true;
+    }
+
     analyzeDeviceToWhere(analysis, queryStatement);
     if (analysis.getWhereExpression() != null
         && analysis.getWhereExpression().equals(ConstantOperand.FALSE)) {

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

@@ -118,7 +118,6 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
@@ -820,7 +819,7 @@ private PlanNode createGroupByTagNode(
       GroupByTimeParameter groupByTimeParameter,
       Ordering scanOrder) {
     Map<List<String>, List<CrossSeriesAggregationDescriptor>> tagValuesToAggregationDescriptors =
-        new HashMap<>();
+        new LinkedHashMap<>();
     for (List<String> tagValues : tagValuesToGroupedTimeseriesOperands.keySet()) {
       LinkedHashMap<Expression, List<Expression>> groupedTimeseriesOperands =
           tagValuesToGroupedTimeseriesOperands.get(tagValues);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/process/GroupByTagNode.java

@@ -38,7 +38,7 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -253,7 +253,7 @@ public static GroupByTagNode deserialize(ByteBuffer byteBuffer) {
     // Tag values to aggregation descriptors.
     int numOfEntries = ReadWriteIOUtils.readInt(byteBuffer);
     Map<List<String>, List<CrossSeriesAggregationDescriptor>> tagValuesToAggregationDescriptors =
-        new HashMap<>();
+        new LinkedHashMap<>();
     while (numOfEntries > 0) {
       List<String> tagValues = ReadWriteIOUtils.readStringList(byteBuffer);
       List<CrossSeriesAggregationDescriptor> aggregationDescriptors = new ArrayList<>();

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/GroupByTagPlanningOrderTest.java

@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.plan.analyze;
+
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.commons.path.PartialPath;
+import org.apache.iotdb.commons.path.PathPatternTree;
+import org.apache.iotdb.commons.queryengine.common.SessionInfo;
+import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
+import org.apache.iotdb.db.queryengine.common.QueryId;
+import org.apache.iotdb.db.queryengine.common.schematree.ClusterSchemaTree;
+import org.apache.iotdb.db.queryengine.plan.expression.Expression;
+import org.apache.iotdb.db.queryengine.plan.parser.StatementGenerator;
+import org.apache.iotdb.db.queryengine.plan.planner.LogicalPlanBuilder;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.GroupByTagNode;
+import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
+
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.write.schema.MeasurementSchema;
+import org.junit.Test;
+
+import java.time.ZoneId;
+import java.time.ZonedDateTime;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+
+public class GroupByTagPlanningOrderTest {
+
+  @Test
+  public void testAnalysisPreservesGroupEncounterOrder() throws Exception {
+    Analysis analysis = analyze(createQueryContext());
+    assertEquals(
+        Arrays.asList(Collections.singletonList("v1"), Collections.singletonList("v0")),
+        new ArrayList<>(analysis.getTagValuesToGroupedTimeseriesOperands().keySet()));
+  }
+
+  @Test
+  public void testLogicalPlanPreservesAnalyzedGroupOrder() throws Exception {
+    MPPQueryContext context = createQueryContext();
+    Analysis analysis = analyze(context);
+    Map<List<String>, LinkedHashMap<Expression, List<Expression>>> groups = new LinkedHashMap<>();
+    for (String value : Arrays.asList("v1", "v0")) {
+      List<String> tagValues = Collections.singletonList(value);
+      groups.put(tagValues, analysis.getTagValuesToGroupedTimeseriesOperands().get(tagValues));
+    }
+
+    GroupByTagNode node =
+        (GroupByTagNode)
+            new LogicalPlanBuilder(analysis, context)
+                .planRawDataSource(analysis.getSourceExpressions(), Ordering.ASC, 0, 0, false)
+                .planGroupByTag(
+                    analysis.getCrossGroupByExpressions(),
+                    analysis.getTagKeys(),
+                    groups,
+                    null,
+                    Ordering.ASC)
+                .getRoot();
+    assertEquals(
+        new ArrayList<>(groups.keySet()),
+        new ArrayList<>(node.getTagValuesToAggregationDescriptors().keySet()));
+  }
+
+  private Analysis analyze(MPPQueryContext context) throws Exception {
+    ClusterSchemaTree testSchemaTree = new ClusterSchemaTree();
+    testSchemaTree.setDatabases(Collections.singleton("root.sg"));
+    for (int i = 1; i <= 2; i++) {
+      testSchemaTree.appendSingleMeasurement(
+          new PartialPath("root.sg.d" + i + ".s1"),
+          new MeasurementSchema("s1", TSDataType.INT32),
+          Collections.singletonMap("key1", i == 1 ? "v1" : "v0"),
+          null,
+          null,
+          false);
+    }
+    FakeSchemaFetcherImpl schemaFetcher =
+        new FakeSchemaFetcherImpl() {
+          @Override
+          public ClusterSchemaTree fetchSchemaWithTags(
+              PathPatternTree patternTree,
+              boolean withTemplate,
+              MPPQueryContext queryContext,
+              boolean canSeeAuditDB) {
+            return testSchemaTree;
+          }
+        };
+    return new Analyzer(context, new FakePartitionFetcherImpl(), schemaFetcher)
+        .analyze(
+            StatementGenerator.createStatement(
+                "select count(s1) from root.sg.* group by tags(key1)",
+                ZonedDateTime.now().getOffset()));
+  }
+
+  private MPPQueryContext createQueryContext() {
+    return new MPPQueryContext(
+        "",
+        new QueryId("test_query"),
+        new SessionInfo(0, "root", ZoneId.systemDefault()),
+        new TEndPoint(),
+        new TEndPoint());
+  }
+}