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

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

commit 29d21941ace81018cb1d9f1ff61ff69dd7e248bf
Author: Beyyes <[email protected]>
AuthorDate: Thu Jun 6 11:55:42 2024 +0800

    set idAndAttributeIndexMap in PruneUnUsedColumns optimize rule
---
 .../plan/relational/planner/RelationPlanner.java           | 12 +-----------
 .../planner/optimizations/PruneUnUsedColumns.java          | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/RelationPlanner.java
 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/RelationPlanner.java
index c09e5d67caa..d39b008a8ca 100644
--- 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/RelationPlanner.java
+++ 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/RelationPlanner.java
@@ -13,7 +13,6 @@
  */
 package org.apache.iotdb.db.queryengine.plan.relational.planner;
 
-import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory;
 import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
 import org.apache.iotdb.db.queryengine.common.QueryId;
 import org.apache.iotdb.db.queryengine.common.SessionInfo;
@@ -41,7 +40,6 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -94,12 +92,11 @@ public class RelationPlanner extends 
AstVisitor<RelationPlan, Void> {
           expansion.getRoot(), expansion.getScope(), 
expansion.getFieldMappings());
     }
 
-    Map<Symbol, Integer> idAndAttributeIndexMap = new HashMap<>();
     Scope scope = analysis.getScope(table);
     ImmutableList.Builder<Symbol> outputSymbolsBuilder = 
ImmutableList.builder();
     ImmutableMap.Builder<Symbol, ColumnSchema> symbolToColumnSchema = 
ImmutableMap.builder();
     Collection<Field> fields = scope.getRelationType().getAllFields();
-    int IDIdx = 0, attributeIdx = 0;
+
     for (Field field : fields) {
       Symbol symbol = symbolAllocator.newSymbol(field);
       outputSymbolsBuilder.add(symbol);
@@ -107,12 +104,6 @@ public class RelationPlanner extends 
AstVisitor<RelationPlan, Void> {
           symbol,
           new ColumnSchema(
               field.getName().get(), field.getType(), field.isHidden(), 
field.getColumnCategory()));
-
-      if (TsTableColumnCategory.ID.equals(field.getColumnCategory())) {
-        idAndAttributeIndexMap.put(symbol, IDIdx++);
-      } else if 
(TsTableColumnCategory.ATTRIBUTE.equals(field.getColumnCategory())) {
-        idAndAttributeIndexMap.put(symbol, attributeIdx++);
-      }
     }
 
     List<Symbol> outputSymbols = outputSymbolsBuilder.build();
@@ -123,7 +114,6 @@ public class RelationPlanner extends 
AstVisitor<RelationPlan, Void> {
             outputSymbols,
             symbolToColumnSchema.build());
 
-    tableScanNode.setIdAndAttributeIndexMap(idAndAttributeIndexMap);
     return new RelationPlan(tableScanNode, scope, outputSymbols);
 
     // Collection<Field> fields = 
analysis.getMaterializedViewStorageTableFields(node);
diff --git 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PruneUnUsedColumns.java
 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PruneUnUsedColumns.java
index d6310c37dc0..c3b0a12c01b 100644
--- 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PruneUnUsedColumns.java
+++ 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PruneUnUsedColumns.java
@@ -14,6 +14,7 @@
 
 package org.apache.iotdb.db.queryengine.plan.relational.planner.optimizations;
 
+import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory;
 import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
 import org.apache.iotdb.db.queryengine.common.SessionInfo;
 import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher;
@@ -119,6 +120,19 @@ public class PruneUnUsedColumns implements 
RelationalPlanOptimizer {
       }
       node.setOutputSymbols(newOutputSymbols);
       node.setAssignments(newAssignments);
+
+      int IDIdx = 0, attributeIdx = 0;
+      Map<Symbol, Integer> idAndAttributeIndexMap = new 
HashMap<>(node.getAssignments().size());
+      for (Symbol symbol : node.getOutputSymbols()) {
+        ColumnSchema columnSchema = node.getAssignments().get(symbol);
+        if (TsTableColumnCategory.ID.equals(columnSchema.getColumnCategory())) 
{
+          idAndAttributeIndexMap.put(symbol, IDIdx++);
+        } else if 
(TsTableColumnCategory.ATTRIBUTE.equals(columnSchema.getColumnCategory())) {
+          idAndAttributeIndexMap.put(symbol, attributeIdx++);
+        }
+      }
+      node.setIdAndAttributeIndexMap(idAndAttributeIndexMap);
+
       return node;
     }
   }

Reply via email to