[VL] Query hangs on destroying memory manager (#4232)

apache · Dec 29, 2023 · fcb31fc · fcb31fc
1 parent 385b252
commit fcb31fc
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 3 deletions.
diff --git a/cpp/velox/compute/VeloxRuntime.cc b/cpp/velox/compute/VeloxRuntime.cc
@@ -161,7 +161,8 @@ std::shared_ptr<Datasource> VeloxRuntime::createDatasource(
     const std::string& filePath,
     MemoryManager* memoryManager,
     std::shared_ptr<arrow::Schema> schema) {
-  auto veloxPool = getAggregateVeloxPool(memoryManager);
+  static std::atomic_uint32_t id{0UL};
+  auto veloxPool = getAggregateVeloxPool(memoryManager)->addAggregateChild("datasource." + std::to_string(id++));
   // Pass a dedicate pool for S3 sink as can't share veloxPool
   // with parquet writer.
   auto s3SinkPool = getLeafVeloxPool(memoryManager);

diff --git a/cpp/velox/memory/VeloxMemoryManager.cc b/cpp/velox/memory/VeloxMemoryManager.cc
@@ -270,8 +270,25 @@ bool VeloxMemoryManager::tryDestructSafe() {
   veloxAggregatePool_.reset();
 
   // Velox memory manager considered safe to destruct when no alive pools.
-  if (veloxMemoryManager_ && veloxMemoryManager_->numPools() != 0) {
-    return false;
+  if (veloxMemoryManager_) {
+    if (veloxMemoryManager_->numPools() > 1) {
+      return false;
+    }
+    if (veloxMemoryManager_->numPools() == 1) {
+      // Assert the pool is spill pool
+      // See https://github.com/facebookincubator/velox/commit/e6f84e8ac9ef6721f527a2d552a13f7e79bdf72e
+      int32_t spillPoolCount = 0;
+      veloxMemoryManager_->testingDefaultRoot().visitChildren([&](velox::memory::MemoryPool* child) -> bool {
+        if (child == veloxMemoryManager_->spillPool()) {
+          spillPoolCount++;
+        }
+        return true;
+      });
+      GLUTEN_CHECK(spillPoolCount == 1, "Illegal pool count state: spillPoolCount: " + std::to_string(spillPoolCount));
+    }
+    if (veloxMemoryManager_->numPools() < 1) {
+      GLUTEN_CHECK(false, "Unreachable code");
+    }
   }
   veloxMemoryManager_.reset();
 

diff --git a/gluten-ut/spark34/src/test/scala/io/glutenproject/utils/velox/VeloxTestSettings.scala b/gluten-ut/spark34/src/test/scala/io/glutenproject/utils/velox/VeloxTestSettings.scala
@@ -769,6 +769,10 @@ class VeloxTestSettings extends BackendTestSettings {
     // error message mismatch is accepted
     .exclude("schema mismatch failure error message for parquet reader")
     .exclude("schema mismatch failure error message for parquet vectorized reader")
+    // [PATH_NOT_FOUND] Path does not exist:
+    // file:/opt/spark331/sql/core/src/test/resources/test-data/timestamp-nanos.parquet
+    // May require for newer spark.test.home
+    .excludeByPrefix("SPARK-40819")
   enableSuite[GlutenParquetThriftCompatibilitySuite]
     // Rewrite for file locating.
     .exclude("Read Parquet file generated by parquet-thrift")