fix: Make TopNRowNumber reclaimable when loading lazy

facebookincubator · Dec 5, 2024 · 859cf1a · 859cf1a
1 parent 37a5ffb
commit 859cf1a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/velox/exec/Operator.h b/velox/exec/Operator.h
@@ -488,7 +488,7 @@ class Operator : public BaseRuntimeStatWriter {
   /// should be called after this.
   virtual void close();
 
-  // Returns true if 'this' never has more output rows than input rows.
+  /// Returns true if 'this' never has more output rows than input rows.
   virtual bool isFilter() const {
     return false;
   }
@@ -832,7 +832,7 @@ std::vector<column_index_t> calculateOutputChannels(
     const RowTypePtr& targetInputType,
     const RowTypePtr& targetOutputType);
 
-// A first operator in a Driver, e.g. table scan or exchange client.
+/// A first operator in a Driver, e.g. table scan or exchange client.
 class SourceOperator : public Operator {
  public:
   SourceOperator(

diff --git a/velox/exec/TopNRowNumber.cpp b/velox/exec/TopNRowNumber.cpp
@@ -183,8 +183,14 @@ void TopNRowNumber::addInput(RowVectorPtr input) {
 
   const auto numInput = input->size();
 
-  for (auto i = 0; i < inputChannels_.size(); ++i) {
-    decodedVectors_[i].decode(*input->childAt(inputChannels_[i]));
+  {
+    // Potential large memory usage site that might trigger arbitration. Make it
+    // reclaimable because at this point it does not break the operator's state
+    // atomicity.
+    ReclaimableSectionGuard guard(this);
+    for (auto i = 0; i < inputChannels_.size(); ++i) {
+      decodedVectors_[i].decode(*input->childAt(inputChannels_[i]));
+    }
   }
 
   if (table_) {