Re-organize projection planner

src/include/planner/query_planner.h

@@ -5,7 +5,6 @@
 #include "binder/query/reading_clause/bound_in_query_call.h"
 #include "join_order_enumerator.h"
 #include "planner/join_order/cardinality_estimator.h"
-#include "projection_planner.h"
 
 namespace kuzu {
 namespace binder {
@@ -30,7 +29,7 @@ class QueryPlanner {
         const storage::RelsStatistics& relsStatistics)
         : catalog{catalog}, cardinalityEstimator{std::make_unique<CardinalityEstimator>(
                                 nodesStatistics, relsStatistics)},
-          joinOrderEnumerator{catalog, this}, projectionPlanner{this} {}
+          joinOrderEnumerator{catalog, this} {}
 
     std::vector<std::unique_ptr<LogicalPlan>> getAllPlans(const BoundStatement& boundStatement);
 
@@ -66,6 +65,16 @@ class QueryPlanner {
     void planSetClause(binder::BoundUpdatingClause& updatingClause, LogicalPlan& plan);
     void planDeleteClause(binder::BoundUpdatingClause& updatingClause, LogicalPlan& plan);
 
+    // Plan projection
+    void planProjectionBody(const binder::BoundProjectionBody& projectionBody,
+        const std::vector<std::unique_ptr<LogicalPlan>>& plans);
+    void planProjectionBody(const binder::BoundProjectionBody& projectionBody, LogicalPlan& plan);
+    void planAggregate(const binder::expression_vector& expressionsToAggregate,
+        const binder::expression_vector& expressionsToGroupBy, LogicalPlan& plan);
+    void planOrderBy(const binder::expression_vector& expressionsToProject,
+        const binder::expression_vector& expressionsToOrderBy, const std::vector<bool>& isAscOrders,
+        LogicalPlan& plan);
+
     // Plan subquery
     void planOptionalMatch(const QueryGraphCollection& queryGraphCollection,
         const expression_vector& predicates, LogicalPlan& leftPlan);
@@ -122,6 +131,16 @@ class QueryPlanner {
     std::vector<std::unique_ptr<BoundSetPropertyInfo>> createLogicalSetPropertyInfo(
         const std::vector<binder::BoundCreateInfo*>& boundCreateInfos);
 
+    // Append projection operators
+    void appendProjection(const binder::expression_vector& expressionsToProject, LogicalPlan& plan);
+    void appendAggregate(const binder::expression_vector& expressionsToGroupBy,
+        const binder::expression_vector& expressionsToAggregate, LogicalPlan& plan);
+    void appendOrderBy(const binder::expression_vector& expressions,
+        const std::vector<bool>& isAscOrders, LogicalPlan& plan);
+    void appendMultiplicityReducer(LogicalPlan& plan);
+    void appendLimit(uint64_t limitNumber, LogicalPlan& plan);
+    void appendSkip(uint64_t skipNumber, LogicalPlan& plan);
+
     std::unique_ptr<LogicalPlan> createUnionPlan(
         std::vector<std::unique_ptr<LogicalPlan>>& childrenPlans, bool isUnionAll);
 
@@ -138,7 +157,6 @@ class QueryPlanner {
     std::unique_ptr<CardinalityEstimator> cardinalityEstimator;
     expression_vector propertiesToScan;
     JoinOrderEnumerator joinOrderEnumerator;
-    ProjectionPlanner projectionPlanner;
 };
 
 } // namespace planner

src/planner/CMakeLists.txt

@@ -7,7 +7,6 @@ add_library(kuzu_planner
         join_order_enumerator.cpp
         join_order_enumerator_context.cpp
         planner.cpp
-        projection_planner.cpp
         query_planner.cpp
         subplans_table.cpp)
 

src/planner/join_order/append_extend.cpp

@@ -181,7 +181,7 @@ void JoinOrderEnumerator::createPathNodePropertyScanPlan(
     auto expressionsToProject = properties;
     expressionsToProject.push_back(recursiveNode->getInternalIDProperty());
     expressionsToProject.push_back(recursiveNode->getLabelExpression());
-    queryPlanner->projectionPlanner.appendProjection(expressionsToProject, plan);
+    queryPlanner->appendProjection(expressionsToProject, plan);
 }
 
 void JoinOrderEnumerator::createPathRelPropertyScanPlan(
@@ -195,7 +195,7 @@ void JoinOrderEnumerator::createPathRelPropertyScanPlan(
     appendNonRecursiveExtend(recursiveNode, nbrNode, recursiveRel, direction, properties, plan);
     auto expressionsToProject = properties;
     expressionsToProject.push_back(recursiveRel->getLabelExpression());
-    queryPlanner->projectionPlanner.appendProjection(expressionsToProject, plan);
+    queryPlanner->appendProjection(expressionsToProject, plan);
 }
 
 void JoinOrderEnumerator::appendNodeLabelFilter(std::shared_ptr<Expression> nodeID,

src/planner/plan/CMakeLists.txt

@@ -1,10 +1,17 @@
 add_library(kuzu_planner_plan_operator
         OBJECT
         append_accumulate.cpp
+        append_aggregate.cpp
         append_create.cpp
         append_cross_product.cpp
         append_delete.cpp
+        append_limit.cpp
+        append_multiplicity_reducer.cpp
+        append_order_by.cpp
+        append_projection.cpp
         append_set.cpp
+        append_skip.cpp
+        plan_projection.cpp
         plan_update.cpp)
 
 set(ALL_OBJECT_FILES

src/planner/plan/append_aggregate.cpp

@@ -0,0 +1,20 @@
+#include "planner/logical_plan/logical_operator/logical_aggregate.h"
+#include "planner/query_planner.h"
+
+namespace kuzu {
+namespace planner {
+
+void QueryPlanner::appendAggregate(const expression_vector& expressionsToGroupBy,
+    const expression_vector& expressionsToAggregate, LogicalPlan& plan) {
+    auto aggregate = make_shared<LogicalAggregate>(
+        expressionsToGroupBy, expressionsToAggregate, plan.getLastOperator());
+    appendFlattens(aggregate->getGroupsPosToFlattenForGroupBy(), plan);
+    aggregate->setChild(0, plan.getLastOperator());
+    appendFlattens(aggregate->getGroupsPosToFlattenForAggregate(), plan);
+    aggregate->setChild(0, plan.getLastOperator());
+    aggregate->computeFactorizedSchema();
+    plan.setLastOperator(std::move(aggregate));
+}
+
+} // namespace planner
+} // namespace kuzu

src/planner/plan/append_limit.cpp

@@ -0,0 +1,17 @@
+#include "planner/logical_plan/logical_operator/logical_limit.h"
+#include "planner/query_planner.h"
+
+namespace kuzu {
+namespace planner {
+
+void QueryPlanner::appendLimit(uint64_t limitNumber, LogicalPlan& plan) {
+    auto limit = make_shared<LogicalLimit>(limitNumber, plan.getLastOperator());
+    appendFlattens(limit->getGroupsPosToFlatten(), plan);
+    limit->setChild(0, plan.getLastOperator());
+    limit->computeFactorizedSchema();
+    plan.setCardinality(limitNumber);
+    plan.setLastOperator(std::move(limit));
+}
+
+} // namespace planner
+} // namespace kuzu

src/planner/plan/append_multiplicity_reducer.cpp

@@ -0,0 +1,14 @@
+#include "planner/logical_plan/logical_operator/logical_multiplcity_reducer.h"
+#include "planner/query_planner.h"
+
+namespace kuzu {
+namespace planner {
+
+void QueryPlanner::appendMultiplicityReducer(LogicalPlan& plan) {
+    auto multiplicityReducer = make_shared<LogicalMultiplicityReducer>(plan.getLastOperator());
+    multiplicityReducer->computeFactorizedSchema();
+    plan.setLastOperator(std::move(multiplicityReducer));
+}
+
+} // namespace planner
+} // namespace kuzu

src/planner/plan/append_order_by.cpp

@@ -0,0 +1,17 @@
+#include "planner/logical_plan/logical_operator/logical_order_by.h"
+#include "planner/query_planner.h"
+
+namespace kuzu {
+namespace planner {
+
+void QueryPlanner::appendOrderBy(
+    const expression_vector& expressions, const std::vector<bool>& isAscOrders, LogicalPlan& plan) {
+    auto orderBy = make_shared<LogicalOrderBy>(expressions, isAscOrders, plan.getLastOperator());
+    appendFlattens(orderBy->getGroupsPosToFlatten(), plan);
+    orderBy->setChild(0, plan.getLastOperator());
+    orderBy->computeFactorizedSchema();
+    plan.setLastOperator(std::move(orderBy));
+}
+
+} // namespace planner
+} // namespace kuzu

src/planner/plan/append_projection.cpp

@@ -0,0 +1,25 @@
+#include "planner/logical_plan/logical_operator/flatten_resolver.h"
+#include "planner/logical_plan/logical_operator/logical_projection.h"
+#include "planner/query_planner.h"
+
+namespace kuzu {
+namespace planner {
+
+void QueryPlanner::appendProjection(
+    const expression_vector& expressionsToProject, LogicalPlan& plan) {
+    for (auto& expression : expressionsToProject) {
+        planSubqueryIfNecessary(expression, plan);
+    }
+    for (auto& expression : expressionsToProject) {
+        auto dependentGroupsPos = plan.getSchema()->getDependentGroupsPos(expression);
+        auto groupsPosToFlatten = factorization::FlattenAllButOne::getGroupsPosToFlatten(
+            dependentGroupsPos, plan.getSchema());
+        appendFlattens(groupsPosToFlatten, plan);
+    }
+    auto projection = make_shared<LogicalProjection>(expressionsToProject, plan.getLastOperator());
+    projection->computeFactorizedSchema();
+    plan.setLastOperator(std::move(projection));
+}
+
+} // namespace planner
+} // namespace kuzu

src/planner/plan/append_skip.cpp

@@ -0,0 +1,17 @@
+#include "planner/logical_plan/logical_operator/logical_skip.h"
+#include "planner/query_planner.h"
+
+namespace kuzu {
+namespace planner {
+
+void QueryPlanner::appendSkip(uint64_t skipNumber, LogicalPlan& plan) {
+    auto skip = make_shared<LogicalSkip>(skipNumber, plan.getLastOperator());
+    appendFlattens(skip->getGroupsPosToFlatten(), plan);
+    skip->setChild(0, plan.getLastOperator());
+    skip->computeFactorizedSchema();
+    plan.setCardinality(plan.getCardinality() - skipNumber);
+    plan.setLastOperator(std::move(skip));
+}
+
+} // namespace planner
+} // namespace kuzu

src/planner/plan/plan_projection.cpp

@@ -0,0 +1,82 @@
+#include "binder/expression_visitor.h"
+#include "planner/query_planner.h"
+
+namespace kuzu {
+namespace planner {
+
+void QueryPlanner::planProjectionBody(const BoundProjectionBody& projectionBody,
+    const std::vector<std::unique_ptr<LogicalPlan>>& plans) {
+    for (auto& plan : plans) {
+        planProjectionBody(projectionBody, *plan);
+    }
+}
+
+void QueryPlanner::planProjectionBody(
+    const BoundProjectionBody& projectionBody, LogicalPlan& plan) {
+    if (plan.isEmpty()) { // e.g. RETURN 1, COUNT(2)
+        appendDummyScan(plan);
+    }
+    // NOTE: As a temporary solution, we rewrite variables in WITH clause as all properties in scope
+    // during planning stage. The purpose is to avoid reading unnecessary properties for WITH.
+    // E.g. MATCH (a) WITH a RETURN a.age -> MATCH (a) WITH a.age RETURN a.age
+    // This rewrite should be removed once we add an optimizer that can remove unnecessary columns.
+    auto expressionsToProject = projectionBody.getProjectionExpressions();
+    auto expressionsToAggregate = projectionBody.getAggregateExpressions();
+    auto expressionsToGroupBy = projectionBody.getGroupByExpressions();
+    if (!expressionsToAggregate.empty()) {
+        planAggregate(expressionsToAggregate, expressionsToGroupBy, plan);
+    }
+    if (projectionBody.hasOrderByExpressions()) {
+        planOrderBy(expressionsToProject, projectionBody.getOrderByExpressions(),
+            projectionBody.getSortingOrders(), plan);
+    }
+    appendProjection(expressionsToProject, plan);
+    if (projectionBody.getIsDistinct()) {
+        appendDistinct(expressionsToProject, plan);
+    }
+    if (projectionBody.hasSkipOrLimit()) {
+        appendMultiplicityReducer(plan);
+        if (projectionBody.hasSkip()) {
+            appendSkip(projectionBody.getSkipNumber(), plan);
+        }
+        if (projectionBody.hasLimit()) {
+            appendLimit(projectionBody.getLimitNumber(), plan);
+        }
+    }
+}
+
+void QueryPlanner::planAggregate(const expression_vector& expressionsToAggregate,
+    const expression_vector& expressionsToGroupBy, LogicalPlan& plan) {
+    assert(!expressionsToAggregate.empty());
+    expression_vector expressionsToProject;
+    for (auto& expressionToAggregate : expressionsToAggregate) {
+        if (ExpressionChildrenCollector::collectChildren(*expressionToAggregate)
+                .empty()) { // skip COUNT(*)
+            continue;
+        }
+        expressionsToProject.push_back(expressionToAggregate->getChild(0));
+    }
+    for (auto& expressionToGroupBy : expressionsToGroupBy) {
+        expressionsToProject.push_back(expressionToGroupBy);
+    }
+    appendProjection(expressionsToProject, plan);
+    appendAggregate(expressionsToGroupBy, expressionsToAggregate, plan);
+}
+
+void QueryPlanner::planOrderBy(const binder::expression_vector& expressionsToProject,
+    const binder::expression_vector& expressionsToOrderBy, const std::vector<bool>& isAscOrders,
+    kuzu::planner::LogicalPlan& plan) {
+    auto expressionsToProjectBeforeOrderBy = expressionsToProject;
+    auto expressionsToProjectSet =
+        expression_set{expressionsToProject.begin(), expressionsToProject.end()};
+    for (auto& expression : expressionsToOrderBy) {
+        if (!expressionsToProjectSet.contains(expression)) {
+            expressionsToProjectBeforeOrderBy.push_back(expression);
+        }
+    }
+    appendProjection(expressionsToProjectBeforeOrderBy, plan);
+    appendOrderBy(expressionsToOrderBy, isAscOrders, plan);
+}
+
+} // namespace planner
+} // namespace kuzu