diff --git a/query_optimizer/physical/Selection.cpp b/query_optimizer/physical/Selection.cpp
index e2fa69e..3307e9e 100644
--- a/query_optimizer/physical/Selection.cpp
+++ b/query_optimizer/physical/Selection.cpp
@@ -42,31 +42,6 @@
 
 namespace E = ::quickstep::optimizer::expressions;
 
-SelectionPtr Selection::Create(
-      const PhysicalPtr &input,
-      const std::vector<E::NamedExpressionPtr> &project_expressions,
-      const E::PredicatePtr &filter_predicate,
-      PartitionSchemeHeader *output_partition_scheme_header) {
-  std::unique_ptr<PartitionSchemeHeader> partition_scheme_header(output_partition_scheme_header);
-
-  if (!partition_scheme_header) {
-    const PartitionSchemeHeader *input_partition_scheme_header = input->getOutputPartitionSchemeHeader();
-    if (input_partition_scheme_header) {
-      unordered_set<E::ExprId> project_expr_ids;
-      for (const E::NamedExpressionPtr &project_expression : project_expressions) {
-        project_expr_ids.insert(project_expression->id());
-      }
-
-      if (input_partition_scheme_header->reusablePartitionScheme(project_expr_ids)) {
-        partition_scheme_header = std::make_unique<PartitionSchemeHeader>(*input_partition_scheme_header);
-      }
-    }
-  }
-
-  return SelectionPtr(
-      new Selection(input, project_expressions, filter_predicate, partition_scheme_header.release()));
-}
-
 PhysicalPtr Selection::copyWithNewChildren(
     const std::vector<PhysicalPtr> &new_children) const {
   DCHECK_EQ(children().size(), new_children.size());
diff --git a/query_optimizer/physical/Selection.hpp b/query_optimizer/physical/Selection.hpp
index 204eb2f..4cbc856 100644
--- a/query_optimizer/physical/Selection.hpp
+++ b/query_optimizer/physical/Selection.hpp
@@ -109,7 +109,10 @@
       const PhysicalPtr &input,
       const std::vector<expressions::NamedExpressionPtr> &project_expressions,
       const expressions::PredicatePtr &filter_predicate,
-      PartitionSchemeHeader *output_partition_scheme_header = nullptr);
+      PartitionSchemeHeader *output_partition_scheme_header = nullptr) {
+    return SelectionPtr(
+        new Selection(input, project_expressions, filter_predicate, output_partition_scheme_header));
+  }
 
   /**
    * @brief Creates a conjunctive predicate with \p filter_predicates
diff --git a/query_optimizer/rules/CollapseSelection.cpp b/query_optimizer/rules/CollapseSelection.cpp
index f92e1b2..fc45ffd 100644
--- a/query_optimizer/rules/CollapseSelection.cpp
+++ b/query_optimizer/rules/CollapseSelection.cpp
@@ -49,7 +49,8 @@
                              { &project_expressions } /* project_expression_lists */);
     return P::Selection::Create(child_selection->input(),
                                 project_expressions,
-                                selection->filter_predicate());
+                                selection->filter_predicate(),
+                                child_selection->input()->cloneOutputPartitionSchemeHeader());
   }
 
   return input;
diff --git a/query_optimizer/rules/ExtractCommonSubexpression.cpp b/query_optimizer/rules/ExtractCommonSubexpression.cpp
index 8b68bd2..4c4d33e 100644
--- a/query_optimizer/rules/ExtractCommonSubexpression.cpp
+++ b/query_optimizer/rules/ExtractCommonSubexpression.cpp
@@ -151,7 +151,8 @@
       if (new_expressions != selection->project_expressions()) {
         return P::Selection::Create(selection->input(),
                                     new_expressions,
-                                    selection->filter_predicate());
+                                    selection->filter_predicate(),
+                                    selection->input()->cloneOutputPartitionSchemeHeader());
       }
       break;
     }
diff --git a/query_optimizer/rules/PushDownLowCostDisjunctivePredicate.cpp b/query_optimizer/rules/PushDownLowCostDisjunctivePredicate.cpp
index 1c71485..69cc299 100644
--- a/query_optimizer/rules/PushDownLowCostDisjunctivePredicate.cpp
+++ b/query_optimizer/rules/PushDownLowCostDisjunctivePredicate.cpp
@@ -195,7 +195,8 @@
     const P::PhysicalPtr selection =
         P::Selection::Create(output,
                              E::ToNamedExpressions(output->getOutputAttributes()),
-                             CreateConjunctive(node_it->second.predicates));
+                             CreateConjunctive(node_it->second.predicates),
+                             output->cloneOutputPartitionSchemeHeader());
 
     // Applicable case 1: The stored relation has small cardinality.
     const bool is_small_cardinality_relation =
diff --git a/query_optimizer/rules/ReorderColumns.cpp b/query_optimizer/rules/ReorderColumns.cpp
index f7e58d5..4783a8d 100644
--- a/query_optimizer/rules/ReorderColumns.cpp
+++ b/query_optimizer/rules/ReorderColumns.cpp
@@ -189,7 +189,8 @@
             std::static_pointer_cast<const P::Selection>(node);
         output = P::Selection::Create(output,
                                       project_expressions,
-                                      old_node->filter_predicate());
+                                      old_node->filter_predicate(),
+                                      output->cloneOutputPartitionSchemeHeader());
         break;
       }
       default:
diff --git a/query_optimizer/strategy/Selection.cpp b/query_optimizer/strategy/Selection.cpp
index 2fbcdec..d738921 100644
--- a/query_optimizer/strategy/Selection.cpp
+++ b/query_optimizer/strategy/Selection.cpp
@@ -119,8 +119,23 @@
     }
   }
 
+  std::unique_ptr<P::PartitionSchemeHeader> physical_output_partition_scheme_header;
+  const P::PartitionSchemeHeader *physical_input_partition_scheme_header =
+      physical_input->getOutputPartitionSchemeHeader();
+  if (physical_input_partition_scheme_header) {
+    std::unordered_set<E::ExprId> project_expr_ids;
+    for (const E::NamedExpressionPtr &project_expression : project_expressions) {
+      project_expr_ids.insert(project_expression->id());
+    }
+
+    if (physical_input_partition_scheme_header->reusablePartitionScheme(project_expr_ids)) {
+      physical_output_partition_scheme_header =
+          std::make_unique<P::PartitionSchemeHeader>(*physical_input_partition_scheme_header);
+    }
+  }
+
   *physical_output = P::Selection::Create(physical_input, project_expressions,
-                                          filter_predicate);
+                                          filter_predicate, physical_output_partition_scheme_header.release());
 }
 
 }  // namespace strategy