Fixed the bug when partition w/ pruned columns.
diff --git a/query_optimizer/ExecutionGenerator.cpp b/query_optimizer/ExecutionGenerator.cpp
index 5ef58a9..555118a 100644
--- a/query_optimizer/ExecutionGenerator.cpp
+++ b/query_optimizer/ExecutionGenerator.cpp
@@ -486,26 +486,8 @@
     const P::PhysicalPtr &physical,
     const CatalogRelation **catalog_relation_output,
     S::InsertDestination *insert_destination_proto) {
-  std::unique_ptr<CatalogRelation> catalog_relation(
-      new CatalogRelation(catalog_database_,
-                          getNewRelationName(),
-                          -1 /* id */,
-                          true /* is_temporary*/));
-  attribute_id aid = 0;
-  for (const E::NamedExpressionPtr &project_expression :
-       physical->getOutputAttributes()) {
-    // The attribute name is simply set to the attribute id to make it distinct.
-    std::unique_ptr<CatalogAttribute> catalog_attribute(
-        new CatalogAttribute(catalog_relation.get(),
-                             std::to_string(aid),
-                             project_expression->getValueType(),
-                             aid,
-                             project_expression->attribute_alias()));
-    attribute_substitution_map_[project_expression->id()] =
-        catalog_attribute.get();
-    catalog_relation->addAttribute(catalog_attribute.release());
-    ++aid;
-  }
+  auto catalog_relation =
+      make_unique<CatalogRelation>(catalog_database_, getNewRelationName(), -1 /* id */, true /* is_temporary*/);
 
   const P::PartitionSchemeHeader *partition_scheme_header = physical->getOutputPartitionSchemeHeader();
   if (partition_scheme_header) {
@@ -514,6 +496,9 @@
       DCHECK(!partition_equivalent_expr_ids.empty());
       const E::ExprId partition_expr_id = *partition_equivalent_expr_ids.begin();
       DCHECK(attribute_substitution_map_.find(partition_expr_id) != attribute_substitution_map_.end());
+      // Use the attribute id from the input relation.
+      // NOTE(zuyu): The following line should be before changing
+      // 'attribute_substitution_map_' with the output attributes.
       output_partition_attr_ids.push_back(attribute_substitution_map_[partition_expr_id]->getID());
     }
 
@@ -544,6 +529,20 @@
     insert_destination_proto->set_insert_destination_type(S::InsertDestinationType::BLOCK_POOL);
   }
 
+  attribute_id aid = 0;
+  for (const E::NamedExpressionPtr &project_expression :
+       physical->getOutputAttributes()) {
+    // The attribute name is simply set to the attribute id to make it distinct.
+    auto catalog_attribute =
+        make_unique<CatalogAttribute>(catalog_relation.get(), std::to_string(aid),
+                                      project_expression->getValueType(), aid,
+                                      project_expression->attribute_alias());
+    attribute_substitution_map_[project_expression->id()] =
+        catalog_attribute.get();
+    catalog_relation->addAttribute(catalog_attribute.release());
+    ++aid;
+  }
+
   *catalog_relation_output = catalog_relation.get();
   const relation_id output_rel_id = catalog_database_->addRelation(
       catalog_relation.release());
diff --git a/query_optimizer/tests/execution_generator/Partition.test b/query_optimizer/tests/execution_generator/Partition.test
index da9b6b8..747b969 100644
--- a/query_optimizer/tests/execution_generator/Partition.test
+++ b/query_optimizer/tests/execution_generator/Partition.test
@@ -297,30 +297,7 @@
 FROM dim_2_hash_partitions, fact
 WHERE dim_2_hash_partitions.id > 20 AND fact.id > 0;
 --
-+-----------+-----------+
-|dim_id     |fact_id    |
-+-----------+-----------+
-|         24|          4|
-|         24|          8|
-|         24|         12|
-|         24|         16|
-|         24|         24|
-|         22|          4|
-|         22|          8|
-|         22|         12|
-|         22|         16|
-|         22|         24|
-|         24|          2|
-|         24|          6|
-|         24|         14|
-|         24|         18|
-|         24|         22|
-|         22|          2|
-|         22|          6|
-|         22|         14|
-|         22|         18|
-|         22|         22|
-+-----------+-----------+
+[same as above]
 ==
 
 SELECT id, int_col