feat(p3): Add top N and optimizer rule (#434)

Add sort, topN

Author: yliang412

src/execution/fmt_impl.cpp

@@ -1,13 +1,15 @@
 #include <type_traits>
-#include "execution/plans/sort_plan.h"
 #include "fmt/format.h"
 #include "fmt/ranges.h"
 
 #include "common/util/string_util.h"
 #include "execution/expressions/abstract_expression.h"
 #include "execution/plans/abstract_plan.h"
 #include "execution/plans/aggregation_plan.h"
+#include "execution/plans/limit_plan.h"
 #include "execution/plans/projection_plan.h"
+#include "execution/plans/sort_plan.h"
+#include "execution/plans/topn_plan.h"
 
 namespace bustub {
 
@@ -40,4 +42,10 @@ auto SortPlanNode::PlanNodeToString() const -> std::string {
   return fmt::format("Sort {{ order_bys={} }}", order_bys_);
 }
 
+auto LimitPlanNode::PlanNodeToString() const -> std::string { return fmt::format("Limit {{ limit={} }}", limit_); }
+
+auto TopNPlanNode::PlanNodeToString() const -> std::string {
+  return fmt::format("TopN {{ n={}, order_bys={}}}", n_, order_bys_);
+}
+
 }  // namespace bustub

src/include/execution/plans/limit_plan.h

@@ -48,7 +48,7 @@ class LimitPlanNode : public AbstractPlanNode {
   BUSTUB_PLAN_NODE_CLONE_WITH_CHILDREN(LimitPlanNode);
 
  protected:
-  auto PlanNodeToString() const -> std::string override { return fmt::format("Limit {{ limit={} }}", limit_); }
+  auto PlanNodeToString() const -> std::string override;
 
  private:
   /** The limit */

src/include/execution/plans/topn_plan.h

@@ -15,7 +15,9 @@
 #include <memory>
 #include <string>
 #include <utility>
+#include <vector>
 
+#include "binder/bound_order_by.h"
 #include "catalog/catalog.h"
 #include "execution/expressions/abstract_expression.h"
 #include "execution/plans/abstract_plan.h"
@@ -33,11 +35,19 @@ class TopNPlanNode : public AbstractPlanNode {
    * @param output The output schema of this topN plan node
    * @param child The child plan node
    */
-  TopNPlanNode(SchemaRef output, AbstractPlanNodeRef child) : AbstractPlanNode(std::move(output), {std::move(child)}) {}
+  TopNPlanNode(SchemaRef output, AbstractPlanNodeRef child,
+               std::vector<std::pair<OrderByType, AbstractExpressionRef>> order_bys, std::size_t n)
+      : AbstractPlanNode(std::move(output), {std::move(child)}), order_bys_(std::move(order_bys)), n_{n} {}
 
   /** @return The type of the plan node */
   auto GetType() const -> PlanType override { return PlanType::TopN; }
 
+  /** @return The N (limit) */
+  auto GetN() const -> size_t { return n_; }
+
+  /** @return Get order by expressions */
+  auto GetOrderBy() const -> const std::vector<std::pair<OrderByType, AbstractExpressionRef>> & { return order_bys_; }
+
   /** @return The child plan node */
   auto GetChildPlan() const -> AbstractPlanNodeRef {
     BUSTUB_ASSERT(GetChildren().size() == 1, "TopN should have exactly one child plan.");
@@ -47,9 +57,11 @@ class TopNPlanNode : public AbstractPlanNode {
   BUSTUB_PLAN_NODE_CLONE_WITH_CHILDREN(TopNPlanNode);
 
  protected:
-  auto PlanNodeToString() const -> std::string override { return fmt::format("TopN {{ }}"); }
+  auto PlanNodeToString() const -> std::string override;
 
  private:
+  std::vector<std::pair<OrderByType, AbstractExpressionRef>> order_bys_;
+  std::size_t n_;
 };
 
 }  // namespace bustub

src/include/optimizer/optimizer.h

@@ -74,6 +74,11 @@ class Optimizer {
   auto MatchIndex(const std::string &table_name, uint32_t index_key_idx)
       -> std::optional<std::tuple<index_oid_t, std::string>>;
 
+  /**
+   * @brief optimize sort + limit as top N
+   */
+  auto OptimizeSortLimitAsTopN(const AbstractPlanNodeRef &plan) -> AbstractPlanNodeRef;
+
   /** Catalog will be used during the planning process. USERS SHOULD ENSURE IT OUTLIVES
    * OPTIMIZER, otherwise it's a dangling reference.
    */

src/optimizer/CMakeLists.txt

@@ -6,7 +6,8 @@ add_library(
     nlj_as_hash_join.cpp
     nlj_as_index_join.cpp
     optimizer.cpp
-    order_by_index_scan.cpp)
+    order_by_index_scan.cpp
+    sort_limit_as_topn.cpp)
 
 set(ALL_OBJECT_FILES
     ${ALL_OBJECT_FILES} $<TARGET_OBJECTS:bustub_optimizer>

src/optimizer/optimizer.cpp

@@ -9,7 +9,8 @@ auto Optimizer::Optimize(const AbstractPlanNodeRef &plan) -> AbstractPlanNodeRef
   auto p3 = OptimizeNLJAsIndexJoin(p2);
   auto p4 = OptimizeNLJAsHashJoin(p3);
   auto p5 = OptimizeOrderByAsIndexScan(p4);
-  return p5;
+  auto p6 = OptimizeSortLimitAsTopN(p5);
+  return p6;
 }
 
 }  // namespace bustub

src/optimizer/sort_limit_as_topn.cpp

@@ -0,0 +1,10 @@
+#include "optimizer/optimizer.h"
+
+namespace bustub {
+
+auto Optimizer::OptimizeSortLimitAsTopN(const AbstractPlanNodeRef &plan) -> AbstractPlanNodeRef {
+  // TODO(student): implement sort + limit -> top N optimizer rule
+  return plan;
+}
+
+}  // namespace bustub