This is an automated email from the ASF dual-hosted git repository.

caogaofei pushed a commit to branch beyyes/TableModel
in repository https://gitbox.apache.org/repos/asf/iotdb.git

commit 1a97aac3dc6e5ad95d55e3c2c40c8c17aeb1ebd1
Author: Beyyes <[email protected]>
AuthorDate: Mon Jun 3 12:02:05 2024 +0800

    add PruneUnUsedColumns optimize rules
---
 .../plan/relational/planner/LogicalPlanner.java    |  4 +--
 .../relational/planner/node/TableScanNode.java     | 20 +++++++----
 ...bleScanColumns.java => PruneUnUsedColumns.java} | 42 +++++++++++++++++++---
 .../plan/relational/analyzer/AnalyzerTest.java     | 18 ++++++++++
 4 files changed, 72 insertions(+), 12 deletions(-)

diff --git 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/LogicalPlanner.java
 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/LogicalPlanner.java
index fb593d4c852..d86c829314d 100644
--- 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/LogicalPlanner.java
+++ 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/LogicalPlanner.java
@@ -29,7 +29,7 @@ import 
org.apache.iotdb.db.queryengine.plan.relational.analyzer.RelationType;
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.OutputNode;
 import 
org.apache.iotdb.db.queryengine.plan.relational.planner.optimizations.IndexScan;
-import 
org.apache.iotdb.db.queryengine.plan.relational.planner.optimizations.PruneTableScanColumns;
+import 
org.apache.iotdb.db.queryengine.plan.relational.planner.optimizations.PruneUnUsedColumns;
 import 
org.apache.iotdb.db.queryengine.plan.relational.planner.optimizations.RelationalPlanOptimizer;
 import 
org.apache.iotdb.db.queryengine.plan.relational.planner.optimizations.RemoveRedundantIdentityProjections;
 import 
org.apache.iotdb.db.queryengine.plan.relational.planner.optimizations.SimplifyExpressions;
@@ -78,7 +78,7 @@ public class LogicalPlanner {
         Arrays.asList(
             new SimplifyExpressions(),
             new RemoveRedundantIdentityProjections(),
-            new PruneTableScanColumns(),
+            new PruneUnUsedColumns(),
             new IndexScan());
   }
 
diff --git 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/TableScanNode.java
 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/TableScanNode.java
index 33dfa3e04bc..33537ec450e 100644
--- 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/TableScanNode.java
+++ 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/node/TableScanNode.java
@@ -45,8 +45,8 @@ public class TableScanNode extends SourceNode {
 
   // db.tablename
   private final String qualifiedTableName;
-  private final List<Symbol> outputSymbols;
-  private final Map<Symbol, ColumnSchema> assignments;
+  private List<Symbol> outputSymbols;
+  private Map<Symbol, ColumnSchema> assignments;
 
   private List<DeviceEntry> deviceEntries;
   private Map<Symbol, Integer> idAndAttributeIndexMap;
@@ -276,6 +276,18 @@ public class TableScanNode extends SourceNode {
     return this.qualifiedTableName;
   }
 
+  public void setOutputSymbols(List<Symbol> outputSymbols) {
+    this.outputSymbols = outputSymbols;
+  }
+
+  public void setAssignments(Map<Symbol, ColumnSchema> assignments) {
+    this.assignments = assignments;
+  }
+
+  public Map<Symbol, ColumnSchema> getAssignments() {
+    return this.assignments;
+  }
+
   public void setDeviceEntries(List<DeviceEntry> deviceEntries) {
     this.deviceEntries = deviceEntries;
   }
@@ -288,10 +300,6 @@ public class TableScanNode extends SourceNode {
     this.idAndAttributeIndexMap = idAndAttributeIndexMap;
   }
 
-  public Map<Symbol, ColumnSchema> getAssignments() {
-    return this.assignments;
-  }
-
   public Ordering getScanOrder() {
     return this.scanOrder;
   }
diff --git 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PruneTableScanColumns.java
 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PruneUnUsedColumns.java
similarity index 68%
rename from 
iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PruneTableScanColumns.java
rename to 
iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PruneUnUsedColumns.java
index cbc18a51465..e896beacd39 100644
--- 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PruneTableScanColumns.java
+++ 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PruneUnUsedColumns.java
@@ -20,9 +20,11 @@ import 
org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNode;
 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.metadata.ColumnSchema;
 import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.Symbol;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.FilterNode;
+import org.apache.iotdb.db.queryengine.plan.relational.planner.node.OutputNode;
 import 
org.apache.iotdb.db.queryengine.plan.relational.planner.node.ProjectNode;
 import 
org.apache.iotdb.db.queryengine.plan.relational.planner.node.TableScanNode;
 import 
org.apache.iotdb.db.queryengine.plan.relational.sql.tree.DefaultTraversalVisitor;
@@ -30,11 +32,24 @@ import 
org.apache.iotdb.db.queryengine.plan.relational.sql.tree.SymbolReference;
 
 import com.google.common.collect.ImmutableList;
 
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
-public class PruneTableScanColumns implements RelationalPlanOptimizer {
+import static org.apache.iotdb.commons.conf.IoTDBConstant.TIME;
+
+/**
+ * Remove unused columns in TableScanNode.
+ *
+ * <p>For example, The output columns of TableScanNode in `select * from 
table1` query are `tag1,
+ * attr1, s1`, but the output columns of TableScanNode in `select s1 from 
table1` query can only be
+ * `s1`.
+ */
+public class PruneUnUsedColumns implements RelationalPlanOptimizer {
+
   @Override
   public PlanNode optimize(
       PlanNode planNode,
@@ -47,6 +62,7 @@ public class PruneTableScanColumns implements 
RelationalPlanOptimizer {
   }
 
   private static class Rewriter extends PlanVisitor<PlanNode, RewriterContext> 
{
+
     @Override
     public PlanNode visitPlan(PlanNode node, RewriterContext context) {
       for (PlanNode child : node.getChildren()) {
@@ -55,9 +71,17 @@ public class PruneTableScanColumns implements 
RelationalPlanOptimizer {
       return node;
     }
 
+    @Override
+    public PlanNode visitOutput(OutputNode node, RewriterContext context) {
+      context.allUsedSymbolSet.addAll(node.getOutputSymbols());
+      node.getChild().accept(this, context);
+      return node;
+    }
+
     @Override
     public PlanNode visitProject(ProjectNode node, RewriterContext context) {
-      context.symbolHashSet.addAll(node.getOutputSymbols());
+      // TODO add project existence judgement
+      context.allUsedSymbolSet.addAll(node.getOutputSymbols());
       node.getChild().accept(this, context);
       return node;
     }
@@ -67,13 +91,23 @@ public class PruneTableScanColumns implements 
RelationalPlanOptimizer {
       ImmutableList.Builder<Symbol> symbolBuilder = ImmutableList.builder();
       new SymbolBuilderVisitor().process(node.getPredicate(), 
ImmutableList.builder());
       List<Symbol> ret = symbolBuilder.build();
-      context.symbolHashSet.addAll(ret);
+      context.allUsedSymbolSet.addAll(ret);
       node.getChild().accept(this, context);
       return node;
     }
 
     @Override
     public PlanNode visitTableScan(TableScanNode node, RewriterContext 
context) {
+      List<Symbol> newOutputSymbols = new ArrayList<>();
+      Map<Symbol, ColumnSchema> newAssignments = new HashMap<>();
+      for (Symbol symbol : node.getOutputSymbols()) {
+        if (TIME.equalsIgnoreCase(symbol.getName()) || 
context.allUsedSymbolSet.contains(symbol)) {
+          newOutputSymbols.add(symbol);
+          newAssignments.put(symbol, node.getAssignments().get(symbol));
+        }
+      }
+      node.setOutputSymbols(newOutputSymbols);
+      node.setAssignments(newAssignments);
       return node;
     }
   }
@@ -89,6 +123,6 @@ public class PruneTableScanColumns implements 
RelationalPlanOptimizer {
   }
 
   private static class RewriterContext {
-    Set<Symbol> symbolHashSet = new HashSet<>();
+    Set<Symbol> allUsedSymbolSet = new HashSet<>();
   }
 }
diff --git 
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/AnalyzerTest.java
 
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/AnalyzerTest.java
index 4eb9a17fb43..42bfe25d562 100644
--- 
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/AnalyzerTest.java
+++ 
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/AnalyzerTest.java
@@ -222,6 +222,24 @@ public class AnalyzerTest {
     assertEquals(4, distributedQueryPlan.getInstances().size());
   }
 
+  @Test
+  public void singleTableProjectTest() throws IoTDBException {
+    // 1. wildcard
+    sql = "SELECT tag1, attr1, s1 FROM table1";
+    actualAnalysis = analyzeSQL(sql, metadata);
+    assertNotNull(actualAnalysis);
+    assertEquals(1, actualAnalysis.getTables().size());
+
+    context = new MPPQueryContext(sql, queryId, sessionInfo, null, null);
+    logicalPlanner =
+        new LogicalPlanner(
+            context, metadata, sessionInfo, getFakePartitionFetcher(), 
WarningCollector.NOOP);
+    logicalQueryPlan = logicalPlanner.plan(actualAnalysis);
+    rootNode = logicalQueryPlan.getRootNode();
+
+    // TODO add TableScanNode outputColumns examination
+  }
+
   public static Analysis analyzeSQL(String sql, Metadata metadata) {
     try {
       SqlParser sqlParser = new SqlParser();

Reply via email to