[fix] admit TopNRowNumber output before non-reclaimable phase (#601)

Author: guhaiyan0221

bolt/common/memory/sparksql/ExecutionMemoryPool.cpp

@@ -408,6 +408,20 @@ uint64_t ExecutionMemoryPool::borrowFromRssWatermarkBytes(
   return it == instance()->borrowFromRssWatermarkBytes_.end() ? 0 : it->second;
 }
 
+uint64_t ExecutionMemoryPool::getConfiguredMemoryPerTask() {
+  if (!inited()) {
+    return 0;
+  }
+
+  auto self = instance();
+  MemoryMutexGuard guard(self->lock_);
+  const auto activeTasks = self->memoryForTask_.size();
+  if (!self->poolExtendSize_.has_value() || activeTasks == 0) {
+    return 0;
+  }
+  return self->poolSize_.value_or(0) / activeTasks;
+}
+
 std::ostream& operator<<(std::ostream& os, const ExecutionMemoryPool& pool) {
   os << "ExecutionMemoryPool(poolSize=" << pool.poolSize_.value_or(0)
      << ", poolExtendSize=" << pool.poolExtendSize_.value_or(0)

bolt/common/memory/sparksql/ExecutionMemoryPool.h

@@ -137,6 +137,8 @@ class ExecutionMemoryPool final
 
   static uint64_t borrowFromRssWatermarkBytes(int64_t taskAttemptId);
 
+  static uint64_t getConfiguredMemoryPerTask();
+
   // For testing only, reset pool size to test new cases with new pool size
   static void testingResetPoolSize(int64_t newSize);
 

bolt/exec/HashBuild.cpp

@@ -1272,6 +1272,8 @@ uint64_t HashBuild::probeAdmissionExtraReservationBytes(
             << succinctBytes(pressure.reclaimWatermarkBytes)
             << ", borrowFromRssWatermarkBytes="
             << succinctBytes(pressure.borrowFromRssWatermarkBytes)
+            << ", configuredTaskMemoryQuotaBytes="
+            << succinctBytes(pressure.configuredTaskMemoryQuotaBytes)
             << ", pressureWatermarkBytes="
             << succinctBytes(pressureWatermarkBytes) << ", details "
             << task->memoryPressureDetails();

bolt/exec/MemoryPressureUtils.cpp

@@ -32,13 +32,18 @@ MemoryPressureSnapshot snapshot(
     uint64_t reclaimWatermarkBytes,
     const std::optional<int64_t>& sparkTaskAttemptId) {
   uint64_t borrowFromRssWatermarkBytes = 0;
+  uint64_t configuredTaskMemoryQuotaBytes = 0;
   if (sparkTaskAttemptId.has_value()) {
     borrowFromRssWatermarkBytes =
         memory::sparksql::ExecutionMemoryPool::borrowFromRssWatermarkBytes(
             *sparkTaskAttemptId);
+    configuredTaskMemoryQuotaBytes =
+        memory::sparksql::ExecutionMemoryPool::getConfiguredMemoryPerTask();
   }
   return MemoryPressureSnapshot{
-      reclaimWatermarkBytes, borrowFromRssWatermarkBytes};
+      reclaimWatermarkBytes,
+      borrowFromRssWatermarkBytes,
+      configuredTaskMemoryQuotaBytes};
 }
 
 std::optional<ScopedMemoryExpansionGuard> maybeScopedDisableMemoryExpansion(

bolt/exec/MemoryPressureUtils.h

@@ -63,10 +63,12 @@ class ScopedMemoryExpansionGuard {
 struct MemoryPressureSnapshot {
   uint64_t reclaimWatermarkBytes{0};
   uint64_t borrowFromRssWatermarkBytes{0};
+  uint64_t configuredTaskMemoryQuotaBytes{0};
 
   uint64_t admissionWatermarkBytes() const {
     if (reclaimWatermarkBytes == 0) {
-      return borrowFromRssWatermarkBytes;
+      return borrowFromRssWatermarkBytes == 0 ? configuredTaskMemoryQuotaBytes
+                                              : borrowFromRssWatermarkBytes;
     }
     if (borrowFromRssWatermarkBytes == 0) {
       return reclaimWatermarkBytes;

bolt/exec/TopNRowNumber.cpp

@@ -293,16 +293,21 @@ void TopNRowNumber::processInputRow(vector_size_t index, TopRows& partition) {
 }
 
 void TopNRowNumber::noMoreInput() {
-  Operator::noMoreInput();
-
   updateEstimatedOutputRowSize();
 
+  if (spiller_ == nullptr) {
+    ensureOutputFits();
+  }
+
+  Operator::noMoreInput();
+
   outputBatchSize_ = outputBatchRows(estimatedOutputRowSize_);
   if (spiller_ != nullptr) {
-    // Spill remaining data to avoid running out of memory while sort-merging
-    // spilled data.
-    spill();
-
+    if (data_->numRows() > 0) {
+      // Spill remaining data to avoid running out of memory while sort-merging
+      // spilled data.
+      spill();
+    }
     BOLT_CHECK_NULL(merge_);
     auto spillPartition = spiller_->finishSpill();
     merge_ = spillPartition.createOrderedReader(pool());
@@ -312,6 +317,59 @@ void TopNRowNumber::noMoreInput() {
   }
 }
 
+void TopNRowNumber::ensureOutputFits() {
+  if (abandonedPartial_ || !spillEnabled() || table_ == nullptr ||
+      data_->numRows() == 0) {
+    return;
+  }
+
+  const auto currentBytes = pool()->currentBytes();
+  const auto* task = operatorCtx_->task().get();
+  const auto pressure = task->memoryPressureSnapshot();
+  const auto pressureWatermarkBytes = pressure.admissionWatermarkBytes();
+  const auto& queryConfig = operatorCtx_->driverCtx()->queryConfig();
+  const uint64_t outputAdmissionLimit =
+      queryConfig.preferredOutputBatchBytes() *
+      queryConfig.outputBatchMemoryReservationMultiple();
+  const auto admissionReservationBytes =
+      (pressureWatermarkBytes != 0 &&
+       currentBytes >
+           pressureWatermarkBytes * queryConfig.memoryPressureWatermarkRatio())
+      ? pressureWatermarkBytes
+      : std::min<uint64_t>(currentBytes / 2, outputAdmissionLimit);
+
+  if (admissionReservationBytes == 0) {
+    return;
+  }
+
+  {
+    Operator::ReclaimableSectionGuard guard(this);
+    const auto scopedDisableMemoryExpansion =
+        task->maybeScopedDisableMemoryExpansion();
+    LOG(INFO) << name() << " ensureOutputFits: try to reserve "
+              << succinctBytes(admissionReservationBytes)
+              << " for output admission, " << pool()->name() << "["
+              << succinctBytes(pool()->currentBytes()) << ", "
+              << succinctBytes(pool()->reservedBytes()) << ", "
+              << succinctBytes(pool()->capacity()) << "]"
+              << ", pressureWatermarkBytes="
+              << succinctBytes(pressureWatermarkBytes)
+              << ", outputAdmissionLimit="
+              << succinctBytes(outputAdmissionLimit)
+              << ", memoryPressure=" << task->memoryPressureDetails();
+    if (pool()->maybeReserve(admissionReservationBytes)) {
+      pool()->release();
+      return;
+    }
+  }
+
+  LOG(WARNING) << name() << " Failed to reserve "
+               << succinctBytes(admissionReservationBytes)
+               << " before producing output for memory pool " << pool()->name()
+               << ", usage: " << succinctBytes(pool()->currentBytes())
+               << ", reservation: " << succinctBytes(pool()->reservedBytes());
+}
+
 void TopNRowNumber::updateEstimatedOutputRowSize() {
   const auto optionalRowSize = data_->estimateRowSize();
   if (!optionalRowSize.has_value()) {

bolt/exec/TopNRowNumber.h

@@ -183,6 +183,11 @@ class TopNRowNumber : public Operator {
   // Called in noMoreInput() and spill().
   void updateEstimatedOutputRowSize();
 
+  // Tries to reserve output admission memory before transitioning to output. If
+  // memory is under pressure, this can trigger arbitration while this operator
+  // is still reclaimable.
+  void ensureOutputFits();
+
   // Sorts, spills and clears all of 'data_'. Clears 'table_'.
   void spill();
 

bolt/exec/tests/TopNRowNumberTest.cpp

@@ -210,6 +210,44 @@ TEST_F(TopNRowNumberTest, preservesOrderAcrossOutputBatches) {
       .assertResults(sql, std::vector<uint32_t>{0, 1});
 }
 
+TEST_F(TopNRowNumberTest, doesNotSpillBeforeOutputWithoutMemoryPressure) {
+  const vector_size_t size = 1'000;
+  auto data = split(
+      makeRowVector({
+          makeFlatVector<int64_t>(size, [](auto row) { return row % 100; }),
+          makeFlatVector<int64_t>(size, [](auto row) { return size - row; }),
+          makeFlatVector<int64_t>(size, [](auto row) { return row; }),
+      }),
+      10);
+
+  createDuckDbTable(data);
+
+  core::PlanNodeId topNRowNumberId;
+  auto plan = PlanBuilder()
+                  .values(data)
+                  .topNRowNumber({"c0"}, {"c1"}, 10, true)
+                  .capturePlanNodeId(topNRowNumberId)
+                  .planNode();
+
+  auto spillDirectory = exec::test::TempDirectoryPath::create();
+  auto task = AssertQueryBuilder(plan, duckDbQueryRunner_)
+                  .config(core::QueryConfig::kSpillEnabled, "true")
+                  .config(core::QueryConfig::kTopNRowNumberSpillEnabled, "true")
+                  .config(core::QueryConfig::kPreferredOutputBatchBytes, "128")
+                  .spillDirectory(spillDirectory->path)
+                  .assertResults(
+                      "SELECT * FROM (SELECT *, row_number() over "
+                      "(partition by c0 order by c1) as rn FROM tmp) "
+                      "WHERE rn <= 10");
+
+  auto taskStats = exec::toPlanStats(task->taskStats());
+  const auto& stats = taskStats.at(topNRowNumberId);
+  ASSERT_EQ(stats.spilledBytes, 0);
+  ASSERT_EQ(stats.spilledRows, 0);
+  ASSERT_EQ(stats.spilledFiles, 0);
+  ASSERT_EQ(stats.spilledPartitions, 0);
+}
+
 TEST_F(TopNRowNumberTest, manyPartitions) {
   const vector_size_t size = 10'000;
   auto data = split(