Advance velox version and use consolidated memory pool

prestodb · Nov 15, 2022 · 05072ca · 05072ca
1 parent 7d5fa52
commit 05072ca
Show file tree

Hide file tree

Showing 9 changed files with 23 additions and 22 deletions.
diff --git a/presto-native-execution/presto_cpp/main/QueryContextManager.cpp b/presto-native-execution/presto_cpp/main/QueryContextManager.cpp
@@ -89,8 +89,8 @@ std::shared_ptr<core::QueryCtx> QueryContextManager::findOrCreateQueryCtx(
   int64_t maxTotalMemoryPerNode = getMaxMemoryPerNode(
       kQueryMaxTotalMemoryPerNode, kDefaultMaxMemoryPerNode);
 
-  auto pool = memory::getProcessDefaultMemoryManager().getRoot().addScopedChild(
-      "query_root");
+  auto pool =
+      memory::getProcessDefaultMemoryManager().getRoot().addChild("query_root");
   pool->setMemoryUsageTracker(velox::memory::MemoryUsageTracker::create(
       maxUserMemoryPerNode, maxSystemMemoryPerNode, maxTotalMemoryPerNode));
 

diff --git a/presto-native-execution/presto_cpp/main/TaskResource.h b/presto-native-execution/presto_cpp/main/TaskResource.h
@@ -23,11 +23,11 @@ class TaskResource {
  public:
   explicit TaskResource(TaskManager& taskManager)
       : taskManager_(taskManager),
-        pool_(velox::memory::getDefaultScopedMemoryPool()) {}
+        pool_(velox::memory::getDefaultMemoryPool()) {}
 
   void registerUris(http::HttpServer& server);
 
-  velox::memory::ScopedMemoryPool* getPool() const {
+  velox::memory::MemoryPool* getPool() const {
     return pool_.get();
   }
 
@@ -65,7 +65,7 @@ class TaskResource {
       const std::vector<std::string>& pathMatch);
 
   TaskManager& taskManager_;
-  std::unique_ptr<velox::memory::ScopedMemoryPool> pool_;
+  std::shared_ptr<velox::memory::MemoryPool> pool_;
 };
 
 } // namespace facebook::presto
diff --git a/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp b/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp
@@ -148,7 +148,7 @@ class TaskManagerTest : public testing::Test {
     connector::registerConnector(hiveConnector);
     dwrf::registerDwrfReaderFactory();
 
-    pool_ = memory::getDefaultScopedMemoryPool();
+    pool_ = memory::getDefaultMemoryPool();
     rowType_ = ROW({"c0", "c1"}, {INTEGER(), VARCHAR()});
 
     taskManager_ = std::make_unique<TaskManager>();
@@ -520,7 +520,7 @@ class TaskManagerTest : public testing::Test {
         fmt::format("{}B", totalMax));
   }
 
-  std::unique_ptr<memory::MemoryPool> pool_;
+  std::shared_ptr<memory::MemoryPool> pool_;
   RowTypePtr rowType_;
   exec::test::DuckDbQueryRunner duckDbQueryRunner_;
   std::unique_ptr<TaskManager> taskManager_;

diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
@@ -1496,16 +1496,17 @@ VeloxQueryPlanConverter::toVeloxQueryPlan(
     groupingSets.emplace_back(std::move(groupingKeys));
   }
 
-  std::map<std::string, core::FieldAccessTypedExprPtr> outputGroupingKeyNames;
+  std::vector<core::GroupIdNode::GroupingKeyInfo> groupingKeys;
+  groupingKeys.reserve(node->groupingColumns.size());
   for (const auto& [output, input] : node->groupingColumns) {
-    outputGroupingKeyNames.emplace(
-        output.name, exprConverter_.toVeloxExpr(input));
+    groupingKeys.emplace_back(core::GroupIdNode::GroupingKeyInfo{
+        output.name, exprConverter_.toVeloxExpr(input)});
   }
 
   return std::make_shared<core::GroupIdNode>(
       node->id,
       std::move(groupingSets),
-      std::move(outputGroupingKeyNames),
+      std::move(groupingKeys),
       toVeloxExprs(node->aggregationArguments),
       node->groupIdVariable.name,
       toVeloxQueryPlan(node->source, tableWriteInfo, taskId));

diff --git a/presto-native-execution/presto_cpp/main/types/tests/PlanConverterTest.cpp b/presto-native-execution/presto_cpp/main/types/tests/PlanConverterTest.cpp
@@ -54,9 +54,9 @@ std::shared_ptr<const core::PlanNode> assertToVeloxQueryPlan(
   std::string fragment = slurp(getDataPath(fileName));
 
   protocol::PlanFragment prestoPlan = json::parse(fragment);
-  auto scopedPool = memory::getDefaultScopedMemoryPool();
+  auto pool = memory::getDefaultMemoryPool();
 
-  VeloxQueryPlanConverter converter(scopedPool.get());
+  VeloxQueryPlanConverter converter(pool.get());
   return converter
       .toVeloxQueryPlan(
           prestoPlan, nullptr, "20201107_130540_00011_wrpkw.1.2.3")

diff --git a/presto-native-execution/presto_cpp/main/types/tests/RowExpressionTest.cpp b/presto-native-execution/presto_cpp/main/types/tests/RowExpressionTest.cpp
@@ -25,7 +25,7 @@ using namespace facebook::velox::core;
 class RowExpressionTest : public ::testing::Test {
  public:
   void SetUp() override {
-    pool_ = memory::getDefaultScopedMemoryPool();
+    pool_ = memory::getDefaultMemoryPool();
     converter_ = std::make_unique<VeloxExprConverter>(pool_.get());
   }
 
@@ -43,7 +43,7 @@ class RowExpressionTest : public ::testing::Test {
     ASSERT_EQ(cexpr->value().toJson(), value);
   }
 
-  std::unique_ptr<memory::MemoryPool> pool_;
+  std::shared_ptr<memory::MemoryPool> pool_;
   std::unique_ptr<VeloxExprConverter> converter_;
 };
 

diff --git a/presto-native-execution/presto_cpp/main/types/tests/ValuesPipeTest.cpp b/presto-native-execution/presto_cpp/main/types/tests/ValuesPipeTest.cpp
@@ -62,8 +62,8 @@ TEST_F(TestValues, valuesRowVector) {
 
   testJsonRoundtrip(j, p);
 
-  auto scopedPool = memory::getDefaultScopedMemoryPool();
-  VeloxQueryPlanConverter converter(scopedPool.get());
+  auto pool = memory::getDefaultMemoryPool();
+  VeloxQueryPlanConverter converter(pool.get());
   auto values = std::dynamic_pointer_cast<const core::ValuesNode>(
       converter.toVeloxQueryPlan(
           std::dynamic_pointer_cast<protocol::PlanNode>(p),
@@ -101,8 +101,8 @@ TEST_F(TestValues, valuesPlan) {
 
   testJsonRoundtrip(j, p);
 
-  auto scopedPool = memory::getDefaultScopedMemoryPool();
-  VeloxQueryPlanConverter converter(scopedPool.get());
+  auto pool = memory::getDefaultMemoryPool();
+  VeloxQueryPlanConverter converter(pool.get());
   auto values = converter.toVeloxQueryPlan(
       std::dynamic_pointer_cast<protocol::OutputNode>(p->root)->source,
       nullptr,

diff --git a/presto-native-execution/presto_cpp/presto_protocol/tests/Base64Test.cpp b/presto-native-execution/presto_cpp/presto_protocol/tests/Base64Test.cpp
@@ -24,10 +24,10 @@ using namespace facebook::velox;
 class Base64Test : public ::testing::Test {
  public:
   void SetUp() override {
-    pool_ = memory::getDefaultScopedMemoryPool();
+    pool_ = memory::getDefaultMemoryPool();
   }
 
-  std::unique_ptr<memory::MemoryPool> pool_;
+  std::shared_ptr<memory::MemoryPool> pool_;
 };
 
 TEST_F(Base64Test, singleInt) {

diff --git a/presto-native-execution/velox b/presto-native-execution/velox