[OPPRO-155] Fix benchmark (facebookincubator#210)

rui-mo · Jun 16, 2022 · 58ede3f · 58ede3f
1 parent 614543e
commit 58ede3f
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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.