perf(cudf): Bound distributed TopN and Grace restore

velox/core/PlanNode.cpp

@@ -2913,17 +2913,23 @@ TopNRowNumberNode::TopNRowNumberNode(
     std::vector<SortOrder> sortingOrders,
     const std::optional<std::string>& rowNumberColumnName,
     int32_t limit,
-    PlanNodePtr source)
+    PlanNodePtr source,
+    bool partialOutput)
     : PlanNode(std::move(id)),
       function_(function),
       partitionKeys_{std::move(partitionKeys)},
       sortingKeys_{std::move(sortingKeys)},
       sortingOrders_{std::move(sortingOrders)},
       limit_{limit},
+      partialOutput_{partialOutput},
       sources_{std::move(source)},
       outputType_{getOptionalRowNumberOutputType(
           sources_[0]->outputType(),
           rowNumberColumnName)} {
+  VELOX_USER_CHECK(
+      !partialOutput_ || !rowNumberColumnName.has_value(),
+      "Partial TopNRowNumber output cannot include a rank column");
+
   VELOX_USER_CHECK_EQ(
       sortingKeys_.size(),
       sortingOrders_.size(),
@@ -2978,6 +2984,7 @@ folly::dynamic TopNRowNumberNode::serialize() const {
     obj["rowNumberColumnName"] = outputType_->names().back();
   }
   obj["limit"] = limit_;
+  obj["partialOutput"] = partialOutput_;
   return obj;
 }
 
@@ -3011,7 +3018,8 @@ PlanNodePtr TopNRowNumberNode::create(
       sortingOrders,
       rowNumberColumnName,
       obj["limit"].asInt(),
-      source);
+      source,
+      obj.count("partialOutput") ? obj["partialOutput"].asBool() : false);
 }
 
 void LocalMergeNode::addDetails(std::stringstream& stream) const {

velox/core/PlanNode.h

@@ -5986,7 +5986,8 @@ class TopNRowNumberNode : public PlanNode {
       std::vector<SortOrder> sortingOrders,
       const std::optional<std::string>& rowNumberColumnName,
       int32_t limit,
-      PlanNodePtr source);
+      PlanNodePtr source,
+      bool partialOutput = false);
 
   class Builder {
    public:
@@ -6004,6 +6005,7 @@ class TopNRowNumberNode : public PlanNode {
       VELOX_CHECK_EQ(other.sources().size(), 1);
       source_ = other.sources()[0];
       function_ = other.rankFunction();
+      partialOutput_ = other.partialOutput();
     }
 
     Builder& id(PlanNodeId id) {
@@ -6047,6 +6049,11 @@ class TopNRowNumberNode : public PlanNode {
       return *this;
     }
 
+    Builder& partialOutput(bool partialOutput) {
+      partialOutput_ = partialOutput;
+      return *this;
+    }
+
     std::shared_ptr<TopNRowNumberNode> build() const {
       VELOX_USER_CHECK(id_.has_value(), "TopNRowNumberNode id is not set");
       VELOX_USER_CHECK(
@@ -6073,7 +6080,8 @@ class TopNRowNumberNode : public PlanNode {
           sortingOrders_.value(),
           rowNumberColumnName_.value(),
           limit_.value(),
-          source_.value());
+          source_.value(),
+          partialOutput_);
     }
 
    private:
@@ -6085,6 +6093,7 @@ class TopNRowNumberNode : public PlanNode {
     std::optional<std::optional<std::string>> rowNumberColumnName_;
     std::optional<int32_t> limit_;
     std::optional<PlanNodePtr> source_;
+    bool partialOutput_{false};
   };
 
   const std::vector<PlanNodePtr>& sources() const override {
@@ -6130,6 +6139,12 @@ class TopNRowNumberNode : public PlanNode {
     return outputType_->size() > sources_[0]->outputType()->size();
   }
 
+  /// Partial TopN may emit one reduced candidate set per input batch. A
+  /// downstream final TopN performs the globally exact reduction.
+  bool partialOutput() const {
+    return partialOutput_;
+  }
+
   std::string_view name() const override {
     return "TopNRowNumber";
   }
@@ -6149,6 +6164,7 @@ class TopNRowNumberNode : public PlanNode {
   const std::vector<SortOrder> sortingOrders_;
 
   const int32_t limit_;
+  const bool partialOutput_;
 
   const std::vector<PlanNodePtr> sources_;
 

velox/core/QueryConfig.h

@@ -348,7 +348,7 @@ class QueryConfig {
       ucxPartitionedOutputBatchRows,
       "cudf.partitioned_output_batch_rows",
       int64_t,
-      10'000,
+      32'000'000,
       "Target rows per CudfPartitionedOutput exchange chunk.")
 
   /// Maximum input rows for one libcudf hash_partition call in UCX output.
@@ -395,7 +395,7 @@ class QueryConfig {
       ucxPartitionedOutputBatchBytes,
       "cudf.partitioned_output_batch_bytes",
       uint64_t,
-      0,
+      128UL << 20,
       "Target bytes per CudfPartitionedOutput exchange chunk.")
 
   VELOX_QUERY_CONFIG(

velox/core/tests/PlanNodeBuilderTest.cpp

@@ -1346,4 +1346,8 @@ TEST_F(PlanNodeBuilderTest, topNRowNumberNode) {
 
   const auto node2 = TopNRowNumberNode::Builder(*node).build();
   verify(node2);
+
+  VELOX_ASSERT_THROW(
+      TopNRowNumberNode::Builder(*node).partialOutput(true).build(),
+      "Partial TopNRowNumber output cannot include a rank column");
 }

velox/core/tests/QueryConfigTest.cpp

@@ -34,6 +34,7 @@ TEST_F(QueryConfigTest, emptyConfig) {
 
   ASSERT_FALSE(config.isLegacyCast());
   EXPECT_EQ(config.maxNumSplitsListenedTo(), 0);
+  EXPECT_EQ(config.ucxPartitionedOutputBatchBytes(), 128UL << 20);
 }
 
 TEST_F(QueryConfigTest, setConfig) {

velox/exec/ExchangeClient.cpp

@@ -107,8 +107,10 @@ void ExchangeClient::close() {
 
 folly::F14FastMap<std::string, RuntimeMetric> ExchangeClient::stats() {
   std::lock_guard<std::mutex> l(queue_->mutex());
-  if (stats_.empty()) {
-    stats_ = collectStatsLocked();
+  if (!closed_) {
+    // Sources are still progressing. Do not freeze the first snapshot: the
+    // exchange operator samples these stats while consuming pages.
+    return collectStatsLocked();
   }
   return stats_;
 }

velox/exec/HashJoinBridge.cpp

@@ -309,12 +309,8 @@ std::optional<HashJoinBridge::HashBuildResult> HashJoinBridge::tableOrFuture(
        restoringSpillShards_.empty()));
   probeStarted_ = true;
   tableSpillFunc_ = nullptr;
-  if (buildResult_.has_value()) {
-    return buildResult_.value();
-  }
-  promises_.emplace_back("HashJoinBridge::tableOrFuture");
-  *future = promises_.back().getSemiFuture();
-  return std::nullopt;
+  return resultOrFutureLocked(
+      buildResult_, future, "HashJoinBridge::tableOrFuture");
 }
 
 void HashJoinBridge::probeFinished(bool restart) {

velox/exec/JoinBridge.h

@@ -15,8 +15,11 @@
  */
 #pragma once
 
+#include "velox/common/base/Exceptions.h"
 #include "velox/common/future/VeloxPromise.h"
 
+#include <optional>
+
 namespace facebook::velox::exec {
 
 class JoinBridge {
@@ -34,6 +37,25 @@ class JoinBridge {
  protected:
   static void notify(std::vector<ContinuePromise> promises);
 
+  /// Returns a bridge result or registers a waiter. The caller must hold
+  /// mutex_. Keeping this transition in JoinBridge gives CPU and accelerator
+  /// bridges the same start/cancel/future contract without constraining the
+  /// result payload type.
+  template <typename Result>
+  std::optional<Result> resultOrFutureLocked(
+      std::optional<Result> result,
+      ContinueFuture* future,
+      const char* waitReason) {
+    VELOX_CHECK(started_);
+    VELOX_CHECK(!cancelled_, "Getting join result after join is aborted");
+    if (result.has_value()) {
+      return result;
+    }
+    promises_.emplace_back(waitReason);
+    *future = promises_.back().getSemiFuture();
+    return std::nullopt;
+  }
+
   std::mutex mutex_;
   bool started_{false};
   std::vector<ContinuePromise> promises_;

velox/exec/Operator.cpp

@@ -738,7 +738,7 @@ bool Operator::MemoryReclaimer::reclaimableBytes(
   if (FOLLY_UNLIKELY(driver == nullptr)) {
     return false;
   }
-  VELOX_CHECK_EQ(pool.name(), op_->pool()->name());
+  VELOX_CHECK_EQ(pool.name(), reclaimPool()->name());
   return op_->reclaimableBytes(reclaimableBytes);
 }
 
@@ -754,7 +754,7 @@ uint64_t Operator::MemoryReclaimer::reclaim(
   if (!op_->canReclaim()) {
     return 0;
   }
-  VELOX_CHECK_EQ(pool->name(), op_->pool()->name());
+  VELOX_CHECK_EQ(pool->name(), reclaimPool()->name());
   VELOX_CHECK(
       !driver->state().isOnThread() || driver->state().suspended() ||
           driver->state().isTerminated,
@@ -810,7 +810,7 @@ void Operator::MemoryReclaimer::abort(
   if (driver == nullptr) {
     return;
   }
-  VELOX_CHECK_EQ(pool->name(), op_->pool()->name());
+  VELOX_CHECK_EQ(pool->name(), reclaimPool()->name());
   VELOX_CHECK(
       !driver->state().isOnThread() || driver->state().suspended() ||
       driver->state().isTerminated);

velox/exec/Operator.h

@@ -582,6 +582,14 @@ class Operator : public BaseRuntimeStatWriter {
       VELOX_CHECK_NOT_NULL(op_);
     }
 
+    /// Returns the pool whose arbitration this reclaimer serves. The default
+    /// preserves the one-operator/one-pool contract. Resource-specific
+    /// operator bases may override this for a leaf in their own parallel pool
+    /// hierarchy without changing the default reclaimer factory API.
+    virtual memory::MemoryPool* reclaimPool() const {
+      return op_->pool();
+    }
+
     // Gets the shared pointer to the associated driver to ensure the liveness
     // of the operator during the memory reclaim operation.
     //

velox/exec/Spill.cpp

@@ -548,62 +548,6 @@ SpillPartition::createOrderedReader(
       spillConfig.readBufferSize, pool, spillStats);
 }
 
-IterableSpillPartitionSet::IterableSpillPartitionSet() {
-  spillPartitionIter_ = spillPartitions_.begin();
-}
-
-void IterableSpillPartitionSet::insert(SpillPartitionSet&& spillPartitionSet) {
-  VELOX_CHECK(
-      !spillPartitionSet.empty(),
-      "Inserted spill partition set must not be empty.");
-
-  const auto parentId = spillPartitionSet.begin()->first.parentId();
-  if (!spillPartitions_.empty()) {
-    VELOX_CHECK(parentId.has_value());
-    VELOX_CHECK(spillPartitionIter_ != spillPartitions_.begin());
-    VELOX_CHECK_EQ(
-        std::prev(spillPartitionIter_)->first,
-        parentId.value(),
-        "Partition set does not have the same parent.");
-    spillPartitions_.erase(std::prev(spillPartitionIter_));
-  } else {
-    VELOX_CHECK(!parentId.has_value());
-  }
-
-  for (const auto& [id, partition] : spillPartitionSet) {
-    VELOX_CHECK_EQ(
-        id.parentId().value_or(SpillPartitionId()),
-        parentId.value_or(SpillPartitionId()));
-    spillPartitions_.emplace(id, std::make_unique<SpillPartition>(*partition));
-  }
-  spillPartitionIter_ = spillPartitions_.find(spillPartitionSet.begin()->first);
-}
-
-bool IterableSpillPartitionSet::hasNext() const {
-  return spillPartitionIter_ != spillPartitions_.end();
-}
-
-SpillPartition IterableSpillPartitionSet::next() {
-  VELOX_CHECK(hasNext(), "No more spill partitions to read.");
-  return *((spillPartitionIter_++)->second);
-}
-
-const SpillPartitionSet& IterableSpillPartitionSet::spillPartitions() const {
-  VELOX_CHECK(
-      !hasNext(),
-      "Spill partitions can only be extracted out after entire set is read.");
-  return spillPartitions_;
-}
-
-void IterableSpillPartitionSet::reset() {
-  spillPartitionIter_ = spillPartitions_.begin();
-}
-
-void IterableSpillPartitionSet::clear() {
-  spillPartitions_.clear();
-  spillPartitionIter_ = spillPartitions_.begin();
-}
-
 uint32_t FileSpillMergeStream::id() const {
   VELOX_CHECK(!closed_);
   return spillFile_->id();