fix

Signed-off-by: Weihao Li <18110526956@163.com>

Author: Wei-hao-Li

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

@@ -25,14 +25,21 @@
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.FillNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GapFillNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.PatternRecognitionNode;
+import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.ProjectNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.SortNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.ValueFillNode;
 import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.WindowNode;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor;
+import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.PruneTableScanColumns;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.DeviceTableScanNode;
 
+import com.google.common.collect.ImmutableSet;
+
 import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
 
 /**
  * <b>Optimization phase:</b> Distributed plan planning.
@@ -42,6 +49,8 @@
  *     SortNode can be eliminated.
  * <li>When order by all IDColumns and time, the SortNode can be eliminated.
  * <li>When StreamSortIndex==OrderBy size()-1, remove this StreamSortNode
+ * <li>After SortNode elimination, visitProject will remove redundant identity ProjectNodes above
+ *     TableScan and pushes column pruning into the scan.
  */
 public class SortElimination implements PlanOptimizer {
 
@@ -64,6 +73,43 @@ public PlanNode visitPlan(PlanNode node, Context context) {
       return newNode;
     }
 
+    @Override
+    public PlanNode visitProject(ProjectNode node, Context context) {
+      Context newContext = new Context();
+      PlanNode child = node.getChild().accept(this, newContext);
+      context.setCannotEliminateSort(newContext.cannotEliminateSort);
+
+      // Remove useless ProjectNode and prune columns of TableScanNode
+      return eliminateProjectOverTableScan(node, child)
+          .orElseGet(() -> node.replaceChildren(Collections.singletonList(child)));
+    }
+
+    private static Optional<PlanNode> eliminateProjectOverTableScan(
+        ProjectNode project, PlanNode child) {
+      if (!(child instanceof DeviceTableScanNode) || !project.isIdentity()) {
+        return Optional.empty();
+      }
+
+      // Notice that SortNode may have been eliminated in TableDistributedPlanGenerator
+      DeviceTableScanNode tableScan = (DeviceTableScanNode) child;
+      int projectOutputsSize = project.getOutputSymbols().size();
+      int scanOutputsSize = tableScan.getOutputSymbols().size();
+      if (projectOutputsSize > scanOutputsSize) {
+        return Optional.empty();
+      }
+
+      List<Symbol> projectOutputs = project.getOutputSymbols();
+      Set<Symbol> scanOutputs = ImmutableSet.copyOf(tableScan.getOutputSymbols());
+      if (!scanOutputs.containsAll(projectOutputs)) {
+        return Optional.empty();
+      }
+
+      if (projectOutputsSize == scanOutputsSize) {
+        return Optional.of(tableScan);
+      }
+      return PruneTableScanColumns.pruneColumns(tableScan, ImmutableSet.copyOf(projectOutputs));
+    }
+
     @Override
     public PlanNode visitSort(SortNode node, Context context) {
       Context newContext = new Context();
@@ -75,9 +121,7 @@ public PlanNode visitSort(SortNode node, Context context) {
           && orderingScheme.getOrderBy().get(0).getName().equals(context.getTimeColumnName())) {
         return child;
       }
-      return context.canEliminateSort() && node.isOrderByAllIdsAndTime()
-          ? child
-          : node.replaceChildren(Collections.singletonList(child));
+      return node.replaceChildren(Collections.singletonList(child));
     }
 
     @Override
@@ -152,6 +196,8 @@ private static class Context {
 
     private String timeColumnName = null;
 
+    private boolean sortEliminated = false;
+
     Context() {}
 
     public void addDeviceEntrySize(int deviceEntrySize) {
@@ -177,5 +223,13 @@ public String getTimeColumnName() {
     public void setTimeColumnName(String timeColumnName) {
       this.timeColumnName = timeColumnName;
     }
+
+    public boolean isSortEliminated() {
+      return sortEliminated;
+    }
+
+    public void markSortEliminated() {
+      this.sortEliminated = true;
+    }
   }
 }

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/SortTest.java

@@ -35,13 +35,16 @@
 import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.sink.IdentitySinkNode;
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata;
+import org.apache.iotdb.db.queryengine.plan.relational.planner.PlanTester;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.SymbolAllocator;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.TableLogicalPlanner;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.distribute.TableDistributedPlanner;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.DeviceTableScanNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.ExchangeNode;
 import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
@@ -58,6 +61,12 @@
 import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestUtils.TEST_MATADATA;
 import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestUtils.assertTableScan;
 import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestUtils.getChildrenNode;
+import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanAssert.assertPlan;
+import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.exchange;
+import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.mergeSort;
+import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.output;
+import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.project;
+import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.tableScan;
 import static org.apache.iotdb.db.queryengine.plan.statement.component.Ordering.ASC;
 import static org.apache.iotdb.db.queryengine.plan.statement.component.Ordering.DESC;
 import static org.junit.Assert.assertEquals;
@@ -767,4 +776,39 @@ public void assertTopKNoFilter(
         expectedPushDownOffset,
         isPushLimitToEachDevice);
   }
+
+  @Test
+  public void singleDeviceOrderByAllIdsAndTimeDescTest() {
+    PlanTester planTester = new PlanTester();
+    planTester.createPlan(
+        "SELECT s1, s2, s3 FROM table1 "
+            + "WHERE time >= 1000 AND time <= 2000 AND tag1='beijing' AND tag2='A1'"
+            + "ORDER BY tag1, tag2, tag3, time DESC");
+    assertPlan(
+        planTester.getFragmentPlan(0),
+        output(
+            tableScan(
+                "testdb.table1",
+                ImmutableList.of("s1", "s2", "s3"),
+                ImmutableSet.of("time", "s1", "s2", "s3"))));
+  }
+
+  @Test
+  public void multiDeviceOrderByAllIdsAndTimeDescTest() {
+    PlanTester planTester = new PlanTester();
+    planTester.createPlan(
+        "SELECT s1, s2, s3 FROM table1 "
+            + "WHERE time >= 1000 AND time <= 2000 AND tag1='beijing' "
+            + "ORDER BY tag1, tag2, tag3, time DESC");
+    assertPlan(
+        planTester.getFragmentPlan(0),
+        output(project(mergeSort(exchange(), exchange(), exchange()))));
+    // Device in multi-region
+    assertPlan(
+        planTester.getFragmentPlan(1),
+        tableScan(
+            "testdb.table1",
+            ImmutableList.of("time", "tag1", "tag2", "tag3", "s1", "s2", "s3"),
+            ImmutableSet.of("time", "tag1", "tag2", "tag3", "s1", "s2", "s3")));
+  }
 }