Fix unreachable ordered rate accumulators

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

@@ -0,0 +1,67 @@
+/*
+ * 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.relational.planner;
+
+import org.apache.iotdb.commons.queryengine.plan.relational.metadata.ColumnSchema;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.Symbol;
+
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+import static org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory.TAG;
+
+/**
+ * Shared checks for whether a set of columns covers every TAG column of a table.
+ *
+ * <p>A scan produces one device per distinct TAG combination, so "every TAG column is covered"
+ * implies that each group of the covering columns belongs to at most one device. Distribution,
+ * grouping parallelization and stream sort all rely on this property, so they must agree on its
+ * definition.
+ */
+public final class TagColumnUtils {
+
+  private TagColumnUtils() {}
+
+  /** Returns the names of all TAG columns in the given table schema. */
+  public static Set<String> tagColumnNames(Map<Symbol, ColumnSchema> tableSchema) {
+    return tagColumnNames(tableSchema.keySet(), tableSchema);
+  }
+
+  /** Returns the names of the TAG columns among the given symbols. */
+  public static Set<String> tagColumnNames(
+      Collection<Symbol> symbols, Map<Symbol, ColumnSchema> columnSchema) {
+    Set<String> tagColumnNames = new HashSet<>();
+    for (Symbol symbol : symbols) {
+      ColumnSchema schema = columnSchema.get(symbol);
+      if (schema != null && schema.getColumnCategory() == TAG) {
+        tagColumnNames.add(schema.getName());
+      }
+    }
+    return tagColumnNames;
+  }
+
+  /** Whether the given column names cover every TAG column of the table. */
+  public static boolean coversAllTagColumns(
+      Map<Symbol, ColumnSchema> tableSchema, Collection<String> coveredColumnNames) {
+    return coveredColumnNames.containsAll(tagColumnNames(tableSchema));
+  }
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java

@@ -105,6 +105,7 @@
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.DeviceEntrySortedMaterializer;
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.spill.SpilledDeviceEntryDataSet;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.SymbolAllocator;
+import org.apache.iotdb.db.queryengine.plan.relational.planner.TagColumnUtils;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationTableScanNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationTreeDeviceViewScanNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AlignedAggregationTreeDeviceViewScanNode;
@@ -2077,7 +2078,9 @@ public List<PlanNode> visitAggregation(AggregationNode node, PlanContext context
     //  push down aggregation if the child of aggregation node only has the union Node
     if (childrenNodes.size() == 1) {
       node.setChild(childrenNodes.get(0));
-      AggregationNode physicalAggregation = withRateFunctionInputOrdering(node, childOrdering);
+      AggregationNode physicalAggregation =
+          withRateFunctionInputOrdering(
+              node, childOrdering, findTimeOrderedDeviceScan(node, childrenNodes.get(0)));
 
       if (childrenNodes.get(0) instanceof UnionNode
           && physicalAggregation.getAggregations().values().stream()
@@ -2246,6 +2249,16 @@ private boolean containsUnorderedCollect(PlanNode node) {
 
   private static AggregationNode withRateFunctionInputOrdering(
       AggregationNode node, OrderingScheme childOrdering) {
+    return withRateFunctionInputOrdering(node, childOrdering, null);
+  }
+
+  /**
+   * @param timeOrderedScan the child's scan when the scan input is already time ordered per group,
+   *     or {@code null} when no such guarantee was established
+   */
+  private static AggregationNode withRateFunctionInputOrdering(
+      AggregationNode node, OrderingScheme childOrdering, DeviceTableScanNode timeOrderedScan) {
+    PlanNode child = node.getChild();
     Map<Symbol, AggregationNode.Aggregation> aggregations = new LinkedHashMap<>();
     node.getAggregations()
         .forEach(
@@ -2259,8 +2272,13 @@ private static AggregationNode withRateFunctionInputOrdering(
                         aggregation.getFilter(),
                         aggregation.getOrderingScheme(),
                         aggregation.getMask(),
-                        isInputOrderedByTimeAscending(
-                            aggregation, node.getStep(), node.getGroupingKeys(), childOrdering))));
+                        (timeOrderedScan != null
+                                && isTimeArgumentNativeTime(aggregation, child, timeOrderedScan))
+                            || isInputOrderedByTimeAscending(
+                                aggregation,
+                                node.getStep(),
+                                node.getGroupingKeys(),
+                                childOrdering))));
     return AggregationNode.builderFrom(node).setAggregations(aggregations).build();
   }
 
@@ -2269,18 +2287,17 @@ static boolean isInputOrderedByTimeAscending(
       AggregationNode.Step step,
       List<Symbol> groupingKeys,
       OrderingScheme childOrdering) {
-    String functionName = aggregation.getResolvedFunction().getSignature().getName();
-    if (step != SINGLE
-        || childOrdering == null
-        || aggregation.getArguments().size() < 2
-        || !(RATE.equalsIgnoreCase(functionName)
-            || INCREASE.equalsIgnoreCase(functionName)
-            || IRATE.equalsIgnoreCase(functionName)
-            || DELTA.equalsIgnoreCase(functionName))) {
+    if (step != SINGLE || childOrdering == null) {
+      return false;
+    }
+    int timeArgumentIndex =
+        getOrderedInputTimeArgumentIndex(
+            aggregation.getResolvedFunction().getSignature().getName());
+    if (timeArgumentIndex < 0 || aggregation.getArguments().size() <= timeArgumentIndex) {
       return false;
     }
 
-    Symbol timeSymbol = Symbol.from(aggregation.getArguments().get(1));
+    Symbol timeSymbol = Symbol.from(aggregation.getArguments().get(timeArgumentIndex));
     List<Symbol> orderBy = childOrdering.getOrderBy();
     int timeIndex = orderBy.indexOf(timeSymbol);
     if (timeIndex < 0 || !childOrdering.getOrdering(timeSymbol).isAscending()) {
@@ -2296,6 +2313,140 @@ static boolean isInputOrderedByTimeAscending(
     return true;
   }
 
+  /**
+   * Returns the child's scan when its input is time ordered per output group, or {@code null}. Only
+   * ordered-input families benefit from the guarantee, so the table schema is not loaded for plain
+   * aggregations such as {@code count} or {@code sum}.
+   */
+  private DeviceTableScanNode findTimeOrderedDeviceScan(AggregationNode node, PlanNode child) {
+    if (node.getAggregations().values().stream()
+        .noneMatch(
+            aggregation ->
+                isOrderedInputFunction(
+                    aggregation.getResolvedFunction().getSignature().getName()))) {
+      return null;
+    }
+    DeviceTableScanNode scan = findOrderedDeviceScan(child);
+    if (scan == null) {
+      return null;
+    }
+    return isDeviceScanTimeOrderedPerGroup(
+            node.getStep(),
+            node.getGroupingKeys(),
+            child,
+            scan,
+            analysis.getTableColumnSchema(scan.getQualifiedObjectName()))
+        ? scan
+        : null;
+  }
+
+  static boolean isDeviceScanTimeOrderedPerGroup(
+      AggregationNode.Step step,
+      List<Symbol> groupingKeys,
+      PlanNode child,
+      DeviceTableScanNode scan,
+      Map<Symbol, ColumnSchema> tableColumnSchema) {
+    if (step != AggregationNode.Step.SINGLE || scan.getScanOrder() != Ordering.ASC) {
+      return false;
+    }
+    int deviceEntryCount = scan.getDeviceEntryCount();
+    if (deviceEntryCount <= 0) {
+      return false;
+    }
+    if (deviceEntryCount == 1) {
+      return true;
+    }
+    if (tableColumnSchema == null) {
+      return false;
+    }
+    Set<String> groupedColumnNames =
+        TagColumnUtils.tagColumnNames(
+            resolveSymbolsToScanLevel(child, groupingKeys), scan.getAssignments());
+    return TagColumnUtils.coversAllTagColumns(tableColumnSchema, groupedColumnNames);
+  }
+
+  static boolean isTimeArgumentNativeTime(
+      AggregationNode.Aggregation aggregation, PlanNode child, DeviceTableScanNode scan) {
+    int timeArgumentIndex =
+        getOrderedInputTimeArgumentIndex(
+            aggregation.getResolvedFunction().getSignature().getName());
+    if (timeArgumentIndex < 0 || aggregation.getArguments().size() <= timeArgumentIndex) {
+      return false;
+    }
+    Expression timeArgument = aggregation.getArguments().get(timeArgumentIndex);
+    if (!(timeArgument instanceof SymbolReference)) {
+      return false;
+    }
+    for (Symbol symbol :
+        resolveSymbolsToScanLevel(child, Collections.singletonList(Symbol.from(timeArgument)))) {
+      if (scan.isTimeColumn(symbol)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  static DeviceTableScanNode findOrderedDeviceScan(PlanNode node) {
+    PlanNode cursor = node;
+    while (cursor instanceof ProjectNode || cursor instanceof FilterNode) {
+      if (cursor.getChildren().size() != 1) {
+        return null;
+      }
+      cursor = cursor.getChildren().get(0);
+    }
+    if (cursor instanceof AggregationTableScanNode
+        || cursor instanceof TreeDeviceViewScanNode
+        || cursor instanceof ExternalTsFileScanNode) {
+      return null;
+    }
+    return cursor instanceof DeviceTableScanNode ? (DeviceTableScanNode) cursor : null;
+  }
+
+  private static Set<Symbol> resolveSymbolsToScanLevel(PlanNode node, Collection<Symbol> symbols) {
+    Set<Symbol> current = new HashSet<>(symbols);
+    PlanNode cursor = node;
+    while (true) {
+      if (cursor instanceof ProjectNode) {
+        Map<Symbol, Expression> assignments = ((ProjectNode) cursor).getAssignments().getMap();
+        Set<Symbol> mapped = new HashSet<>();
+        for (Symbol symbol : current) {
+          Expression expression = assignments.get(symbol);
+          if (expression instanceof SymbolReference) {
+            mapped.add(Symbol.from(expression));
+          }
+        }
+        current = mapped;
+        cursor = cursor.getChildren().get(0);
+      } else if (cursor instanceof FilterNode) {
+        cursor = cursor.getChildren().get(0);
+      } else {
+        return current;
+      }
+    }
+  }
+
+  /**
+   * Returns the position of the time argument in an aggregation's argument list, or {@code -1} when
+   * the function's accumulators do not require time-ordered input. Both ordering checks must use
+   * this helper: a new family (for example the integral family, whose time argument is {@code
+   * integral(value, time_unit, time)}) only needs to be registered here, and its time argument
+   * index then applies to the ordering-based and the scan-based check alike.
+   */
+  private static int getOrderedInputTimeArgumentIndex(String functionName) {
+    // RATE, INCREASE, IRATE and DELTA take (value, time, ...).
+    if (RATE.equalsIgnoreCase(functionName)
+        || INCREASE.equalsIgnoreCase(functionName)
+        || IRATE.equalsIgnoreCase(functionName)
+        || DELTA.equalsIgnoreCase(functionName)) {
+      return 1;
+    }
+    return -1;
+  }
+
+  private static boolean isOrderedInputFunction(String functionName) {
+    return getOrderedInputTimeArgumentIndex(functionName) >= 0;
+  }
+
   private boolean prefixMatched(OrderingScheme childOrdering, List<Symbol> preGroupedSymbols) {
     List<Symbol> orderKeys = childOrdering.getOrderBy();
     if (orderKeys.size() < preGroupedSymbols.size()) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java

@@ -37,6 +37,7 @@
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor;
 import org.apache.iotdb.db.queryengine.plan.relational.analyzer.Analysis;
+import org.apache.iotdb.db.queryengine.plan.relational.planner.TagColumnUtils;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationTableScanNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.DeviceTableScanNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.ExternalTsFileAggregationScanNode;
@@ -276,11 +277,6 @@ public PlanNode visitDeviceTableScan(DeviceTableScanNode node, Context context)
           tableColumnSchema = analysis.getTableColumnSchema(node.getQualifiedObjectName());
         }
         //  check there are no field in sortKey and all tags in sortKey
-        Set<Symbol> tagSymbols =
-            tableColumnSchema.entrySet().stream()
-                .filter(entry -> entry.getValue().getColumnCategory() == TAG)
-                .map(Map.Entry::getKey)
-                .collect(Collectors.toSet());
         for (int i = 0; i < context.partitionKeyCount; i++) {
           Symbol symbol = sortKey.getOrderBy().get(i);
           if (!tableColumnSchema.containsKey(symbol)) {
@@ -289,7 +285,6 @@ public PlanNode visitDeviceTableScan(DeviceTableScanNode node, Context context)
           }
           switch (tableColumnSchema.get(symbol).getColumnCategory()) {
             case TAG:
-              tagSymbols.remove(symbol);
               break;
             case ATTRIBUTE:
               // If all tags in partition key, attributes must be the same in one partition.
@@ -299,11 +294,13 @@ public PlanNode visitDeviceTableScan(DeviceTableScanNode node, Context context)
               return node;
           }
         }
-        if (!tagSymbols.isEmpty()) {
-          context.canOptimized = TO_SORT;
-        } else {
-          context.canOptimized = ENABLE_PARALLEL;
-        }
+        Set<String> tagColumnNames =
+            TagColumnUtils.tagColumnNames(
+                sortKey.getOrderBy().subList(0, context.partitionKeyCount), tableColumnSchema);
+        context.canOptimized =
+            TagColumnUtils.coversAllTagColumns(tableColumnSchema, tagColumnNames)
+                ? ENABLE_PARALLEL
+                : TO_SORT;
       }
       return node;
     }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/TransformSortToStreamSort.java

@@ -33,6 +33,7 @@
 import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor;
 import org.apache.iotdb.db.queryengine.plan.relational.analyzer.Analysis;
+import org.apache.iotdb.db.queryengine.plan.relational.planner.TagColumnUtils;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationTableScanNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.CteScanNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.DeviceTableScanNode;
@@ -224,20 +225,9 @@ public static boolean isOrderByAllIdsAndTime(
       Map<Symbol, ColumnSchema> nodeColumnSchema,
       OrderingScheme orderingScheme,
       int streamSortIndex) {
-    int tagCount = 0;
-    for (ColumnSchema columnSchema : tableColumnSchema.values()) {
-      if (columnSchema.getColumnCategory() == TsTableColumnCategory.TAG) {
-        tagCount++;
-      }
-    }
-
-    for (Symbol orderBy : orderingScheme.getOrderBy()) {
-      ColumnSchema columnSchema = nodeColumnSchema.get(orderBy);
-      if (columnSchema != null && columnSchema.getColumnCategory() == TsTableColumnCategory.TAG) {
-        tagCount--;
-      }
-    }
-    return tagCount == 0
+    Set<String> orderedTagColumnNames =
+        TagColumnUtils.tagColumnNames(orderingScheme.getOrderBy(), nodeColumnSchema);
+    return TagColumnUtils.coversAllTagColumns(tableColumnSchema, orderedTagColumnNames)
         && (orderingScheme.getOrderings().size() == streamSortIndex + 1
             || isTimeColumn(
                 orderingScheme.getOrderBy().get(streamSortIndex + 1), tableColumnSchema));