From 58ede3fbf06e65c3e09b16a3ff7ed033226d3547 Mon Sep 17 00:00:00 2001
From: Rui Mo <rui.mo@intel.com>
Date: Thu, 16 Jun 2022 15:04:17 +0800
Subject: [PATCH] [OPPRO-155] Fix benchmark (#210)

---
 cpp/velox/benchmarks/BenchmarkUtils.cc    | 4 ++--
 cpp/velox/benchmarks/BenchmarkUtils.h     | 3 ++-
 cpp/velox/benchmarks/HashJoinBenchmark.cc | 5 +++--
 cpp/velox/benchmarks/QueryBenchmark.cc    | 3 ++-
 4 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/cpp/velox/benchmarks/BenchmarkUtils.cc b/cpp/velox/benchmarks/BenchmarkUtils.cc
index 70898d7ff365..ad398a6e5f90 100644
--- a/cpp/velox/benchmarks/BenchmarkUtils.cc
+++ b/cpp/velox/benchmarks/BenchmarkUtils.cc
@@ -34,9 +34,9 @@ std::string getExampleFilePath(const std::string& fileName) {
                                                   "data/" + fileName);
 }
 
-void InitVeloxBackend() {
+void InitVeloxBackend(facebook::velox::memory::MemoryPool* pool) {
   gluten::SetBackendFactory(
-      [] { return std::make_shared<::velox::compute::VeloxPlanConverter>(); });
+      [=] { return std::make_shared<::velox::compute::VeloxPlanConverter>(pool); });
   auto veloxInitializer = std::make_shared<::velox::compute::VeloxInitializer>();
 }
 
diff --git a/cpp/velox/benchmarks/BenchmarkUtils.h b/cpp/velox/benchmarks/BenchmarkUtils.h
index 0cfad3075898..ef364a3ca3c0 100644
--- a/cpp/velox/benchmarks/BenchmarkUtils.h
+++ b/cpp/velox/benchmarks/BenchmarkUtils.h
@@ -18,9 +18,10 @@
 #pragma once
 
 #include "compute/protobuf_utils.h"
+#include "velox/common/memory/Memory.h"
 
 /// Initilize the Velox backend.
-void InitVeloxBackend();
+void InitVeloxBackend(facebook::velox::memory::MemoryPool* pool);
 
 /// Get the location of a file in this project.
 std::string getExampleFilePath(const std::string& fileName);
diff --git a/cpp/velox/benchmarks/HashJoinBenchmark.cc b/cpp/velox/benchmarks/HashJoinBenchmark.cc
index 628402b155d4..b49dbc398462 100644
--- a/cpp/velox/benchmarks/HashJoinBenchmark.cc
+++ b/cpp/velox/benchmarks/HashJoinBenchmark.cc
@@ -106,9 +106,10 @@ auto BM_HashJoinExample = [](::benchmark::State& state) {
   """).write.format("parquet").save("file:///path/to/bm_part")
  */
 int main(int argc, char** argv) {
-  InitVeloxBackend();
+  std::unique_ptr<facebook::velox::memory::MemoryPool> veloxPool =
+      facebook::velox::memory::getDefaultScopedMemoryPool();
+  InitVeloxBackend(veloxPool.get());
   ::benchmark::Initialize(&argc, argv);
-
   if (argc < 3) {
     std::cout << "Running example." << std::endl;
     ::benchmark::RegisterBenchmark("hash_join_example", BM_HashJoinExample);
diff --git a/cpp/velox/benchmarks/QueryBenchmark.cc b/cpp/velox/benchmarks/QueryBenchmark.cc
index 67931e83b3de..5f00ad5af06d 100644
--- a/cpp/velox/benchmarks/QueryBenchmark.cc
+++ b/cpp/velox/benchmarks/QueryBenchmark.cc
@@ -50,7 +50,8 @@ auto BM = [](::benchmark::State& state, const std::string& datasetPath,
 };
 
 int main(int argc, char** argv) {
-  InitVeloxBackend();
+  std::unique_ptr<memory::MemoryPool> veloxPool = memory::getDefaultScopedMemoryPool();
+  InitVeloxBackend(veloxPool.get());
   ::benchmark::Initialize(&argc, argv);
   // Threads cannot work well, use ThreadRange instead.
   // The multi-thread performance is not correct.