feat(planner): support count(*)

src/include/binder/bound_expression.h

@@ -98,9 +98,6 @@ struct fmt::formatter<bustub::ExpressionType> : formatter<string_view> {
       case bustub::ExpressionType::ALIAS:
         name = "Alias";
         break;
-      default:
-        name = "Unknown";
-        break;
     }
     return formatter<string_view>::format(name, ctx);
   }

src/include/execution/executors/aggregation_executor.h

@@ -47,6 +47,7 @@ class SimpleAggregationHashTable {
     std::vector<Value> values{};
     for (const auto &agg_type : agg_types_) {
       switch (agg_type) {
+        case AggregationType::CountStarAggregate:
         case AggregationType::CountAggregate:
         case AggregationType::SumAggregate:
           // Count/Sum starts at zero.
@@ -73,21 +74,11 @@ class SimpleAggregationHashTable {
   void CombineAggregateValues(AggregateValue *result, const AggregateValue &input) {
     for (uint32_t i = 0; i < agg_exprs_.size(); i++) {
       switch (agg_types_[i]) {
+        case AggregationType::CountStarAggregate:
         case AggregationType::CountAggregate:
-          // Count increases by one.
-          result->aggregates_[i] = result->aggregates_[i].Add(ValueFactory::GetIntegerValue(1));
-          break;
         case AggregationType::SumAggregate:
-          // Sum increases by addition.
-          result->aggregates_[i] = result->aggregates_[i].Add(input.aggregates_[i]);
-          break;
         case AggregationType::MinAggregate:
-          // Min is just the min.
-          result->aggregates_[i] = result->aggregates_[i].Min(input.aggregates_[i]);
-          break;
         case AggregationType::MaxAggregate:
-          // Max is just the max.
-          result->aggregates_[i] = result->aggregates_[i].Max(input.aggregates_[i]);
           break;
       }
     }

src/include/execution/plans/aggregation_plan.h

@@ -26,7 +26,7 @@
 namespace bustub {
 
 /** AggregationType enumerates all the possible aggregation functions in our system */
-enum class AggregationType { CountAggregate, SumAggregate, MinAggregate, MaxAggregate };
+enum class AggregationType { CountStarAggregate, CountAggregate, SumAggregate, MinAggregate, MaxAggregate };
 
 /**
  * AggregationPlanNode represents the various SQL aggregation functions.
@@ -147,6 +147,9 @@ struct fmt::formatter<bustub::AggregationType> : formatter<std::string> {
     using bustub::AggregationType;
     std::string name = "unknown";
     switch (c) {
+      case AggregationType::CountStarAggregate:
+        name = "count_star";
+        break;
       case AggregationType::CountAggregate:
         name = "count";
         break;

src/include/planner/planner.h

@@ -117,7 +117,13 @@ class Planner {
   auto PlanSelectAgg(const SelectStatement &statement, AbstractPlanNodeRef child) -> AbstractPlanNodeRef;
 
   auto PlanAggCall(const BoundAggCall &agg_call, const std::vector<AbstractPlanNodeRef> &children)
-      -> std::tuple<AggregationType, AbstractExpressionRef>;
+      -> std::tuple<AggregationType, std::vector<AbstractExpressionRef>>;
+
+  auto GetAggCallFromFactory(const std::string &func_name, std::vector<AbstractExpressionRef> args)
+      -> std::tuple<AggregationType, std::vector<AbstractExpressionRef>>;
+
+  auto GetBinaryExpressionFromFactory(const std::string &op_name, AbstractExpressionRef left,
+                                      AbstractExpressionRef right) -> AbstractExpressionRef;
 
   auto PlanInsert(const InsertStatement &statement) -> AbstractPlanNodeRef;
 

src/planner/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_library(
   bustub_planner
   OBJECT
+  expression_factory.cpp
   plan_aggregation.cpp
   plan_expression.cpp
   plan_insert.cpp

src/planner/expression_factory.cpp

@@ -0,0 +1,74 @@
+#include "binder/bound_expression.h"
+#include "execution/expressions/abstract_expression.h"
+#include "execution/expressions/arithmetic_expression.h"
+#include "execution/expressions/column_value_expression.h"
+#include "execution/expressions/comparison_expression.h"
+#include "execution/expressions/constant_value_expression.h"
+#include "execution/expressions/logic_expression.h"
+#include "planner/planner.h"
+
+namespace bustub {
+// NOLINTNEXTLINE - weird error on clang-tidy.
+auto Planner::GetAggCallFromFactory(const std::string &func_name, std::vector<AbstractExpressionRef> args)
+    -> std::tuple<AggregationType, std::vector<AbstractExpressionRef>> {
+  if (args.empty()) {
+    if (func_name == "count_star") {
+      return {AggregationType::CountStarAggregate, {}};
+    }
+  }
+  if (args.size() == 1) {
+    auto expr = std::move(args[0]);
+    if (func_name == "min") {
+      return {AggregationType::MinAggregate, {std::move(expr)}};
+    }
+    if (func_name == "max") {
+      return {AggregationType::MaxAggregate, {std::move(expr)}};
+    }
+    if (func_name == "sum") {
+      return {AggregationType::SumAggregate, {std::move(expr)}};
+    }
+    if (func_name == "count") {
+      return {AggregationType::CountAggregate, {std::move(expr)}};
+    }
+  }
+  throw Exception(fmt::format("unsupported agg_call {} with {} args", func_name, args.size()));
+}
+
+auto Planner::GetBinaryExpressionFromFactory(const std::string &op_name, AbstractExpressionRef left,
+                                             AbstractExpressionRef right) -> AbstractExpressionRef {
+  if (op_name == "=" || op_name == "==") {
+    return std::make_shared<ComparisonExpression>(std::move(left), std::move(right), ComparisonType::Equal);
+  }
+  if (op_name == "!=" || op_name == "<>") {
+    return std::make_shared<ComparisonExpression>(std::move(left), std::move(right), ComparisonType::NotEqual);
+  }
+  if (op_name == "<") {
+    return std::make_shared<ComparisonExpression>(std::move(left), std::move(right), ComparisonType::LessThan);
+  }
+  if (op_name == "<=") {
+    return std::make_shared<ComparisonExpression>(std::move(left), std::move(right), ComparisonType::LessThanOrEqual);
+  }
+  if (op_name == ">") {
+    return std::make_shared<ComparisonExpression>(std::move(left), std::move(right), ComparisonType::GreaterThan);
+  }
+  if (op_name == ">=") {
+    return std::make_shared<ComparisonExpression>(std::move(left), std::move(right),
+                                                  ComparisonType::GreaterThanOrEqual);
+  }
+  if (op_name == "+") {
+    return std::make_shared<ArithmeticExpression>(std::move(left), std::move(right), ArithmeticType::Plus);
+  }
+  if (op_name == "-") {
+    return std::make_shared<ArithmeticExpression>(std::move(left), std::move(right), ArithmeticType::Minus);
+  }
+  if (op_name == "and") {
+    return std::make_shared<LogicExpression>(std::move(left), std::move(right), LogicType::And);
+  }
+  if (op_name == "or") {
+    return std::make_shared<LogicExpression>(std::move(left), std::move(right), LogicType::Or);
+  }
+
+  throw Exception(fmt::format("binary op {} not supported in planner yet", op_name));
+}
+
+}  // namespace bustub

src/planner/plan_aggregation.cpp

@@ -12,45 +12,36 @@
 #include "common/util/string_util.h"
 #include "execution/expressions/abstract_expression.h"
 #include "execution/expressions/column_value_expression.h"
+#include "execution/expressions/constant_value_expression.h"
 #include "execution/plans/abstract_plan.h"
 #include "execution/plans/aggregation_plan.h"
 #include "execution/plans/filter_plan.h"
 #include "execution/plans/projection_plan.h"
 #include "fmt/format.h"
 #include "planner/planner.h"
 #include "type/type_id.h"
+#include "type/value_factory.h"
 
 namespace bustub {
 
 auto Planner::PlanAggCall(const BoundAggCall &agg_call, const std::vector<AbstractPlanNodeRef> &children)
-    -> std::tuple<AggregationType, AbstractExpressionRef> {
-  if (agg_call.args_.size() != 1) {
-    throw NotImplementedException("only agg call of one arg is supported for now");
-  }
+    -> std::tuple<AggregationType, std::vector<AbstractExpressionRef>> {
   if (agg_call.is_distinct_) {
     throw NotImplementedException("distinct agg is not implemented yet");
   }
 
-  AbstractExpressionRef expr = nullptr;
+  std::vector<AbstractExpressionRef> exprs;
+
   {
     // Create a new context that doesn't allow aggregation calls.
     auto guard = NewContext();
-    auto [_, ret] = PlanExpression(*agg_call.args_[0], children);
-    expr = std::move(ret);
-  }
-  if (agg_call.func_name_ == "min") {
-    return {AggregationType::MinAggregate, std::move(expr)};
-  }
-  if (agg_call.func_name_ == "max") {
-    return {AggregationType::MaxAggregate, std::move(expr)};
-  }
-  if (agg_call.func_name_ == "sum") {
-    return {AggregationType::SumAggregate, std::move(expr)};
-  }
-  if (agg_call.func_name_ == "count") {
-    return {AggregationType::CountAggregate, std::move(expr)};
+    for (const auto &arg : agg_call.args_) {
+      auto [_, ret] = PlanExpression(*arg, children);
+      exprs.emplace_back(std::move(ret));
+    }
   }
-  throw Exception(fmt::format("unsupported agg_call {}", agg_call.func_name_));
+
+  return GetAggCallFromFactory(agg_call.func_name_, std::move(exprs));
 }
 
 // TODO(chi): clang-tidy on macOS will suggest changing it to const reference. Looks like a bug.
@@ -121,8 +112,17 @@ auto Planner::PlanSelectAgg(const SelectStatement &statement, AbstractPlanNodeRe
       throw NotImplementedException("alias for agg call is not supported for now");
     }
     const auto &agg_call = dynamic_cast<const BoundAggCall &>(*item);
-    auto [agg_type, expr] = PlanAggCall(agg_call, {child});
-    input_exprs.push_back(std::move(expr));
+    auto [agg_type, exprs] = PlanAggCall(agg_call, {child});
+    if (exprs.size() > 1) {
+      throw bustub::NotImplementedException("only agg call of zero/one arg is supported");
+    }
+    if (exprs.empty()) {
+      // Rewrite count(*) into count(1)
+      input_exprs.emplace_back(std::make_shared<ConstantValueExpression>(ValueFactory::GetIntegerValue(1)));
+    } else {
+      input_exprs.emplace_back(std::move(exprs[0]));
+    }
+
     agg_types.push_back(agg_type);
     output_col_names.emplace_back(fmt::format("agg#{}", term_idx));
     ctx_.expr_in_agg_.emplace_back(

src/planner/plan_expression.cpp

@@ -11,12 +11,8 @@
 #include "common/exception.h"
 #include "common/macros.h"
 #include "common/util/string_util.h"
-#include "execution/expressions/abstract_expression.h"
-#include "execution/expressions/arithmetic_expression.h"
 #include "execution/expressions/column_value_expression.h"
-#include "execution/expressions/comparison_expression.h"
 #include "execution/expressions/constant_value_expression.h"
-#include "execution/expressions/logic_expression.h"
 #include "execution/plans/abstract_plan.h"
 #include "fmt/format.h"
 #include "planner/planner.h"
@@ -28,40 +24,7 @@ auto Planner::PlanBinaryOp(const BoundBinaryOp &expr, const std::vector<Abstract
   auto [_1, left] = PlanExpression(*expr.larg_, children);
   auto [_2, right] = PlanExpression(*expr.rarg_, children);
   const auto &op_name = expr.op_name_;
-
-  if (op_name == "=" || op_name == "==") {
-    return std::make_shared<ComparisonExpression>(std::move(left), std::move(right), ComparisonType::Equal);
-  }
-  if (op_name == "!=" || op_name == "<>") {
-    return std::make_shared<ComparisonExpression>(std::move(left), std::move(right), ComparisonType::NotEqual);
-  }
-  if (op_name == "<") {
-    return std::make_shared<ComparisonExpression>(std::move(left), std::move(right), ComparisonType::LessThan);
-  }
-  if (op_name == "<=") {
-    return std::make_shared<ComparisonExpression>(std::move(left), std::move(right), ComparisonType::LessThanOrEqual);
-  }
-  if (op_name == ">") {
-    return std::make_shared<ComparisonExpression>(std::move(left), std::move(right), ComparisonType::GreaterThan);
-  }
-  if (op_name == ">=") {
-    return std::make_shared<ComparisonExpression>(std::move(left), std::move(right),
-                                                  ComparisonType::GreaterThanOrEqual);
-  }
-  if (op_name == "+") {
-    return std::make_shared<ArithmeticExpression>(std::move(left), std::move(right), ArithmeticType::Plus);
-  }
-  if (op_name == "-") {
-    return std::make_shared<ArithmeticExpression>(std::move(left), std::move(right), ArithmeticType::Minus);
-  }
-  if (op_name == "and") {
-    return std::make_shared<LogicExpression>(std::move(left), std::move(right), LogicType::And);
-  }
-  if (op_name == "or") {
-    return std::make_shared<LogicExpression>(std::move(left), std::move(right), LogicType::Or);
-  }
-
-  throw Exception(fmt::format("binary op {} not supported in planner yet", op_name));
+  return GetBinaryExpressionFromFactory(op_name, std::move(left), std::move(right));
 }
 
 auto Planner::PlanColumnRef(const BoundColumnRef &expr, const std::vector<AbstractPlanNodeRef> &children)