Profile pipeline refactor

src/binder/bind/bind_explain.cpp

@@ -8,7 +8,8 @@ namespace binder {
 std::unique_ptr<BoundStatement> Binder::bindExplain(const parser::Statement& statement) {
     auto& explainStatement = (parser::ExplainStatement&)statement;
     auto boundStatementToExplain = bind(*explainStatement.getStatementToExplain());
-    return std::make_unique<BoundExplain>(std::move(boundStatementToExplain));
+    return std::make_unique<BoundExplain>(
+        std::move(boundStatementToExplain), explainStatement.getExplainType());
 }
 
 } // namespace binder

src/include/binder/bound_explain.h

@@ -1,21 +1,26 @@
 #pragma once
 
 #include "binder/bound_statement.h"
+#include "common/explain_type.h"
 
 namespace kuzu {
 namespace binder {
 
 class BoundExplain : public BoundStatement {
 public:
-    explicit BoundExplain(std::unique_ptr<BoundStatement> statementToExplain)
+    explicit BoundExplain(
+        std::unique_ptr<BoundStatement> statementToExplain, common::ExplainType explainType)
         : BoundStatement{common::StatementType::EXPLAIN,
               BoundStatementResult::createSingleStringColumnResult()},
-          statementToExplain{std::move(statementToExplain)} {}
+          statementToExplain{std::move(statementToExplain)}, explainType{explainType} {}
 
     inline BoundStatement* getStatementToExplain() const { return statementToExplain.get(); }
 
+    inline common::ExplainType getExplainType() const { return explainType; }
+
 private:
     std::unique_ptr<BoundStatement> statementToExplain;
+    common::ExplainType explainType;
 };
 
 } // namespace binder

src/include/common/explain_type.h

@@ -0,0 +1,14 @@
+#pragma once
+
+#include <cstdint>
+
+namespace kuzu {
+namespace common {
+
+enum class ExplainType : uint8_t {
+    PROFILE = 0,
+    PHYSICAL_PLAN = 1,
+};
+
+} // namespace common
+} // namespace kuzu

src/include/main/plan_printer.h

@@ -88,7 +88,8 @@ class OpProfileTree {
 class PlanPrinter {
 
 public:
-    PlanPrinter(processor::PhysicalPlan* physicalPlan, std::unique_ptr<common::Profiler> profiler);
+    PlanPrinter(processor::PhysicalPlan* physicalPlan, common::Profiler* profiler)
+        : physicalPlan{physicalPlan}, profiler{profiler} {};
 
     nlohmann::json printPlanToJson();
 
@@ -104,7 +105,7 @@ class PlanPrinter {
 
 private:
     processor::PhysicalPlan* physicalPlan;
-    std::unique_ptr<common::Profiler> profiler;
+    common::Profiler* profiler;
 };
 
 } // namespace main

src/include/main/prepared_statement.h

@@ -1,5 +1,10 @@
 #pragma once
 
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
 #include "common/api.h"
 #include "kuzu_fwd.h"
 #include "query_summary.h"
@@ -43,10 +48,13 @@ class PreparedStatement {
     inline std::unordered_map<std::string, std::shared_ptr<common::Value>> getParameterMap() {
         return parameterMap;
     }
+
     ~PreparedStatement();
 
 private:
-    common::StatementType statementType;
+    bool isProfile();
+
+private:
     bool success = true;
     bool readOnly = false;
     std::string errMsg;

src/include/main/query_summary.h

@@ -1,9 +1,7 @@
 #pragma once
 
 #include "common/api.h"
-#include "json_fwd.hpp"
 #include "kuzu_fwd.h"
-#include "plan_printer.h"
 
 namespace kuzu {
 namespace main {
@@ -13,9 +11,7 @@ namespace main {
  */
 struct PreparedSummary {
     double compilingTime = 0;
-    // Only used for printing by shell.
-    bool isExplain = false;
-    bool isProfile = false;
+    common::StatementType statementType;
 };
 
 /**
@@ -34,29 +30,17 @@ class QuerySummary {
      * @return query execution time in milliseconds.
      */
     KUZU_API double getExecutionTime() const;
-    bool getIsProfile() const;
-    std::ostringstream& getPlanAsOstream();
-    /**
-     * @return physical plan for query in string format.
-     */
-    KUZU_API std::string getPlan();
+
     void setPreparedSummary(PreparedSummary preparedSummary_);
 
     /**
      * @return true if the query is executed with EXPLAIN.
      */
-    inline bool isExplain() const { return preparedSummary.isExplain; }
-
-private:
-    nlohmann::json& printPlanToJson();
+    bool isExplain() const;
 
 private:
     double executionTime = 0;
     PreparedSummary preparedSummary;
-    // Remove these two field once we have refactored the profiler using the existing pipeline
-    // design.
-    std::unique_ptr<nlohmann::json> planInJson;
-    std::ostringstream planInOstream;
 };
 
 } // namespace main

src/include/parser/explain_statement.h

@@ -1,20 +1,25 @@
 #pragma once
 
+#include "common/explain_type.h"
 #include "parser/statement.h"
 
 namespace kuzu {
 namespace parser {
 
 class ExplainStatement : public Statement {
 public:
-    explicit ExplainStatement(std::unique_ptr<Statement> statementToExplain)
-        : Statement{common::StatementType::EXPLAIN}, statementToExplain{
-                                                         std::move(statementToExplain)} {}
+    explicit ExplainStatement(
+        std::unique_ptr<Statement> statementToExplain, common::ExplainType explainType)
+        : Statement{common::StatementType::EXPLAIN},
+          statementToExplain{std::move(statementToExplain)}, explainType{explainType} {}
 
     inline Statement* getStatementToExplain() const { return statementToExplain.get(); }
 
+    inline common::ExplainType getExplainType() const { return explainType; }
+
 private:
     std::unique_ptr<Statement> statementToExplain;
+    common::ExplainType explainType;
 };
 
 } // namespace parser

src/include/planner/logical_plan/logical_operator/logical_explain.h

@@ -1,17 +1,22 @@
 #pragma once
 
 #include "base_logical_operator.h"
+#include "common/explain_type.h"
 
 namespace kuzu {
 namespace planner {
 
 class LogicalExplain : public LogicalOperator {
 public:
     LogicalExplain(std::shared_ptr<LogicalOperator> child,
-        std::shared_ptr<binder::Expression> outputExpression)
-        : LogicalOperator{LogicalOperatorType::EXPLAIN, child}, outputExpression{
-                                                                    std::move(outputExpression)} {}
-
+        std::shared_ptr<binder::Expression> outputExpression, common::ExplainType explainType,
+        binder::expression_vector outputExpressionsToExplain)
+        : LogicalOperator{LogicalOperatorType::EXPLAIN, child}, outputExpression{std::move(
+                                                                    outputExpression)},
+          explainType{explainType}, outputExpressionsToExplain{
+                                        std::move(outputExpressionsToExplain)} {}
+
+    void computeSchema();
     void computeFactorizedSchema() override;
     void computeFlatSchema() override;
 
@@ -21,12 +26,21 @@ class LogicalExplain : public LogicalOperator {
 
     inline std::string getExpressionsForPrinting() const override { return "Explain"; }
 
+    inline common::ExplainType getExplainType() const { return explainType; }
+
+    inline binder::expression_vector getOutputExpressionsToExplain() const {
+        return outputExpressionsToExplain;
+    }
+
     inline std::unique_ptr<LogicalOperator> copy() override {
-        return std::make_unique<LogicalExplain>(children[0], outputExpression);
+        return std::make_unique<LogicalExplain>(
+            children[0], outputExpression, explainType, outputExpressionsToExplain);
     }
 
-protected:
+private:
     std::shared_ptr<binder::Expression> outputExpression;
+    common::ExplainType explainType;
+    binder::expression_vector outputExpressionsToExplain;
 };
 
 } // namespace planner

src/include/planner/logical_plan/logical_plan.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "planner/logical_plan/logical_operator/base_logical_operator.h"
+#include "planner/logical_plan/logical_operator/logical_explain.h"
 
 namespace kuzu {
 namespace planner {
@@ -29,6 +30,12 @@ class LogicalPlan {
 
     inline std::string toString() const { return lastOperator->toString(); }
 
+    inline bool isProfile() const {
+        return lastOperator->getOperatorType() == LogicalOperatorType::EXPLAIN &&
+               reinterpret_cast<LogicalExplain*>(lastOperator.get())->getExplainType() ==
+                   common::ExplainType::PROFILE;
+    }
+
     std::unique_ptr<LogicalPlan> shallowCopy() const;
 
     std::unique_ptr<LogicalPlan> deepCopy() const;

src/include/processor/mapper/plan_mapper.h

@@ -139,6 +139,13 @@ class PlanMapper {
     static std::vector<DataPos> getExpressionsDataPos(
         const binder::expression_vector& expressions, const planner::Schema& schema);
 
+    std::unique_ptr<PhysicalOperator> appendResultCollectorIfNotCopy(
+        std::unique_ptr<PhysicalOperator> lastOperator,
+        binder::expression_vector expressionsToCollect, planner::Schema* schema);
+
+    static void setPhysicalPlanIfProfile(
+        planner::LogicalPlan* logicalPlan, PhysicalPlan* physicalPlan);
+
 public:
     storage::StorageManager& storageManager;
     storage::MemoryManager* memoryManager;

src/include/processor/operator/physical_operator.h

@@ -41,6 +41,7 @@ enum class PhysicalOperatorType : uint8_t {
     MULTIPLICITY_REDUCER,
     PATH_PROPERTY_PROBE,
     PROJECTION,
+    PROFILE,
     RECURSIVE_JOIN,
     RENAME_PROPERTY,
     RENAME_TABLE,

src/include/processor/operator/profile.h

@@ -0,0 +1,44 @@
+#pragma once
+
+#include "processor/operator/physical_operator.h"
+
+namespace kuzu {
+namespace processor {
+
+struct ProfileInfo {
+    PhysicalPlan* physicalPlan;
+};
+
+struct ProfileLocalState {
+    bool hasExecuted = false;
+};
+
+class Profile : public PhysicalOperator {
+public:
+    Profile(DataPos outputPos, ProfileInfo info, ProfileLocalState localState, uint32_t id,
+        const std::string& paramsString, std::unique_ptr<PhysicalOperator> child)
+        : PhysicalOperator{PhysicalOperatorType::PROFILE, std::move(child), id, paramsString},
+          outputPos{outputPos}, info{std::move(info)}, localState{std::move(localState)} {}
+
+    inline bool isSource() const override { return true; }
+
+    inline void setPhysicalPlan(PhysicalPlan* physicalPlan) { info.physicalPlan = physicalPlan; }
+
+    void initLocalStateInternal(ResultSet* resultSet, ExecutionContext* context) override;
+
+    bool getNextTuplesInternal(ExecutionContext* context) override;
+
+    std::unique_ptr<PhysicalOperator> clone() override {
+        return std::make_unique<Profile>(
+            outputPos, info, localState, id, paramsString, children[0]->clone());
+    }
+
+private:
+    DataPos outputPos;
+    ProfileInfo info;
+    ProfileLocalState localState;
+    common::ValueVector* outputVector;
+};
+
+} // namespace processor
+} // namespace kuzu

src/main/connection.cpp

@@ -2,10 +2,10 @@
 
 #include "binder/binder.h"
 #include "binder/visitor/statement_read_write_analyzer.h"
-#include "json.hpp"
 #include "main/database.h"
 #include "main/plan_printer.h"
 #include "optimizer/optimizer.h"
+#include "parser/explain_statement.h"
 #include "parser/parser.h"
 #include "planner/logical_plan/logical_plan_util.h"
 #include "planner/planner.h"
@@ -160,13 +160,10 @@ std::unique_ptr<PreparedStatement> Connection::prepareNoLock(
     try {
         // parsing
         auto statement = Parser::parseQuery(query);
-        preparedStatement->preparedSummary.isProfile = statement->isProfile();
-        preparedStatement->preparedSummary.isExplain =
-            statement->getStatementType() == StatementType::EXPLAIN;
         // binding
         auto binder = Binder(*database->catalog);
         auto boundStatement = binder.bind(*statement);
-        preparedStatement->statementType = boundStatement->getStatementType();
+        preparedStatement->preparedSummary.statementType = boundStatement->getStatementType();
         preparedStatement->readOnly =
             binder::StatementReadWriteAnalyzer().isReadOnly(*boundStatement);
         preparedStatement->parameterMap = binder.getParameterMap();
@@ -355,7 +352,7 @@ std::unique_ptr<QueryResult> Connection::executeAndAutoCommitIfNecessaryNoLock(
     auto executionContext =
         std::make_unique<ExecutionContext>(clientContext->numThreadsForExecution, profiler.get(),
             database->memoryManager.get(), database->bufferManager.get(), clientContext.get());
-    profiler->enabled = preparedStatement->preparedSummary.isProfile;
+    profiler->enabled = preparedStatement->isProfile();
     auto executingTimer = TimeMetric(true /* enable */);
     executingTimer.start();
     std::shared_ptr<FactorizedTable> resultFT;
@@ -372,10 +369,6 @@ std::unique_ptr<QueryResult> Connection::executeAndAutoCommitIfNecessaryNoLock(
     queryResult->initResultTableAndIterator(std::move(resultFT),
         preparedStatement->statementResult->getColumns(),
         preparedStatement->statementResult->getExpressionsToCollectPerColumn());
-    auto planPrinter = std::make_unique<PlanPrinter>(physicalPlan.get(), std::move(profiler));
-    queryResult->querySummary->planInJson =
-        std::make_unique<nlohmann::json>(planPrinter->printPlanToJson());
-    queryResult->querySummary->planInOstream = planPrinter->printPlanToOstream();
     return queryResult;
 }
 

src/main/plan_printer.cpp

@@ -297,10 +297,6 @@ uint32_t OpProfileTree::calculateRowHeight(uint32_t rowIdx) const {
     return height + 2;
 }
 
-PlanPrinter::PlanPrinter(
-    processor::PhysicalPlan* physicalPlan, std::unique_ptr<common::Profiler> profiler)
-    : physicalPlan{physicalPlan}, profiler{std::move(profiler)} {}
-
 nlohmann::json PlanPrinter::printPlanToJson() {
     return toJson(physicalPlan->lastOperator.get(), *profiler);
 }

src/main/prepared_statement.cpp

@@ -8,7 +8,7 @@ namespace kuzu {
 namespace main {
 
 bool PreparedStatement::allowActiveTransaction() const {
-    return !common::StatementTypeUtils::isDDLOrCopyCSV(statementType);
+    return !common::StatementTypeUtils::isDDLOrCopyCSV(preparedSummary.statementType);
 }
 
 bool PreparedStatement::isSuccess() const {
@@ -27,6 +27,10 @@ binder::expression_vector PreparedStatement::getExpressionsToCollect() {
     return statementResult->getExpressionsToCollect();
 }
 
+bool PreparedStatement::isProfile() {
+    return logicalPlans[0]->isProfile();
+}
+
 PreparedStatement::~PreparedStatement() = default;
 
 } // namespace main