Remove jo connection

src/include/main/connection.h

@@ -16,6 +16,7 @@ class Connection {
     friend class kuzu::testing::ApiTest;
     friend class kuzu::testing::BaseGraphTest;
     friend class kuzu::testing::TestHelper;
+    friend class kuzu::benchmark::Benchmark;
 
 public:
     /**
@@ -129,6 +130,7 @@ class Connection {
     ConnectionTransactionMode getTransactionMode();
     void setTransactionModeNoLock(ConnectionTransactionMode newTransactionMode);
 
+    std::unique_ptr<QueryResult> query(const std::string& query, const std::string& encodedJoin);
     // Note: This is only added for testing recovery algorithms in unit tests. Do not use
     // this otherwise.
     void commitButSkipCheckpointingForTestingRecovery();
@@ -151,8 +153,8 @@ class Connection {
 
     std::unique_ptr<QueryResult> queryResultWithError(std::string& errMsg);
 
-    std::unique_ptr<PreparedStatement> prepareNoLock(
-        const std::string& query, bool enumerateAllPlans = false);
+    std::unique_ptr<PreparedStatement> prepareNoLock(const std::string& query,
+        bool enumerateAllPlans = false, std::string joinOrder = std::string{});
 
     template<typename T, typename... Args>
     std::unique_ptr<QueryResult> executeWithParams(PreparedStatement* preparedStatement,

src/include/main/database.h

@@ -38,7 +38,6 @@ KUZU_API struct SystemConfig {
 class Database {
     friend class EmbeddedShell;
     friend class Connection;
-    friend class JOConnection;
     friend class kuzu::testing::BaseGraphTest;
 
 public:

src/include/main/prepared_statement.h

@@ -13,7 +13,6 @@ namespace main {
  */
 class PreparedStatement {
     friend class Connection;
-    friend class JOConnection;
     friend class kuzu::testing::TestHelper;
     friend class kuzu::transaction::TinySnbDDLTest;
     friend class kuzu::transaction::TinySnbCopyCSVTransactionTest;

src/main/connection.cpp

@@ -6,6 +6,7 @@
 #include "main/plan_printer.h"
 #include "optimizer/optimizer.h"
 #include "parser/parser.h"
+#include "planner/logical_plan/logical_plan_util.h"
 #include "planner/planner.h"
 #include "processor/mapper/plan_mapper.h"
 #include "processor/processor.h"
@@ -73,8 +74,14 @@ std::unique_ptr<PreparedStatement> Connection::prepare(const std::string& query)
 
 std::unique_ptr<QueryResult> Connection::query(const std::string& query) {
     lock_t lck{mtx};
-    std::unique_ptr<PreparedStatement> preparedStatement;
-    preparedStatement = prepareNoLock(query);
+    auto preparedStatement = prepareNoLock(query);
+    return executeAndAutoCommitIfNecessaryNoLock(preparedStatement.get());
+}
+
+std::unique_ptr<QueryResult> Connection::query(
+    const std::string& query, const std::string& encodedJoin) {
+    lock_t lck{mtx};
+    auto preparedStatement = prepareNoLock(query, true /* enumerate all plans */, encodedJoin);
     return executeAndAutoCommitIfNecessaryNoLock(preparedStatement.get());
 }
 
@@ -138,7 +145,7 @@ void Connection::rollbackIfNecessaryNoLock() {
 }
 
 std::unique_ptr<PreparedStatement> Connection::prepareNoLock(
-    const std::string& query, bool enumerateAllPlans) {
+    const std::string& query, bool enumerateAllPlans, std::string encodedJoin) {
     auto preparedStatement = std::make_unique<PreparedStatement>();
     if (query.empty()) {
         preparedStatement->success = false;
@@ -173,10 +180,24 @@ std::unique_ptr<PreparedStatement> Connection::prepareNoLock(
             plans.push_back(Planner::getBestPlan(
                 *database->catalog, nodeStatistics, relStatistics, *boundStatement));
         }
+        // optimizing
         for (auto& plan : plans) {
             optimizer::Optimizer::optimize(plan.get());
         }
-        preparedStatement->logicalPlans = std::move(plans);
+        if (!encodedJoin.empty()) {
+            std::unique_ptr<LogicalPlan> match;
+            for (auto& plan : plans) {
+                if (LogicalPlanUtil::encodeJoin(*plan) == encodedJoin) {
+                    match = std::move(plan);
+                }
+            }
+            if (match == nullptr) {
+                throw ConnectionException("Cannot find a plan matching " + encodedJoin);
+            }
+            preparedStatement->logicalPlans.push_back(std::move(match));
+        } else {
+            preparedStatement->logicalPlans = std::move(plans);
+        }
     } catch (std::exception& exception) {
         preparedStatement->success = false;
         preparedStatement->errMsg = exception.what();

tools/CMakeLists.txt

@@ -1,6 +1,5 @@
 add_subdirectory(shell)
 add_subdirectory(python_api)
 if(${BUILD_BENCHMARK})
-    add_subdirectory(join_order_pick)
     add_subdirectory(benchmark)
 endif()

tools/benchmark/CMakeLists.txt

@@ -7,13 +7,12 @@ set(CMAKE_CXX_STANDARD 20)
 include_directories(
         ../../src/include
         ./include
-        ../../test/test_helper/include
-        ../../tools/join_order_pick/include)
+        ../../test/include/test_helper)
 link_directories(../../build/release/src)
 
 add_executable(kuzu_benchmark
         benchmark.cpp
         benchmark_runner.cpp
         main.cpp)
 
-target_link_libraries(kuzu_benchmark kuzu join_order_pick test_helper)
+target_link_libraries(kuzu_benchmark kuzu test_helper)

tools/benchmark/benchmark.cpp

@@ -13,7 +13,7 @@ namespace benchmark {
 
 Benchmark::Benchmark(const std::string& benchmarkPath, Database* database, BenchmarkConfig& config)
     : config{config} {
-    conn = std::make_unique<JOConnection>(database);
+    conn = std::make_unique<Connection>(database);
     conn->setMaxNumThreadForExec(config.numThreads);
     loadBenchmark(benchmarkPath);
 }

tools/benchmark/include/benchmark.h

@@ -1,7 +1,6 @@
 #pragma once
 
 #include "benchmark_config.h"
-#include "jo_connection.h"
 #include "main/kuzu.h"
 
 namespace kuzu {
@@ -26,7 +25,7 @@ class Benchmark {
 
 public:
     BenchmarkConfig& config;
-    std::unique_ptr<main::JOConnection> conn;
+    std::unique_ptr<main::Connection> conn;
     std::string name;
     std::string query;
     std::vector<std::string> expectedOutput;