diff --git a/query_optimizer/PhysicalGenerator.cpp b/query_optimizer/PhysicalGenerator.cpp
index 6932b30..865cd11 100644
--- a/query_optimizer/PhysicalGenerator.cpp
+++ b/query_optimizer/PhysicalGenerator.cpp
@@ -173,6 +173,7 @@
   // set output PartitionSchemeHeader in a Physical Plan node, when needed.
   if (FLAGS_use_partition_rule) {
     rules.push_back(std::make_unique<Partition>(optimizer_context_));
+    rules.push_back(std::make_unique<PruneColumns>());
   }
 
   // NOTE(jianqiao): Adding rules after InjectJoinFilters (or AttachLIPFilters)