feat: Add PlanNode::requiresSingleThread() API (#16753)

Summary:
Pull Request resolved: #16753

Add a virtual method to PlanNode that returns true if the node requires
single-threaded execution (maxDrivers = 1). Override in ValuesNode,
ArrowStreamNode, final TopNNode, final LimitNode, final OrderByNode,
LocalMergeNode, MergeExchangeNode, MergeJoinNode, TableWriteMergeNode,
LocalPartitionNode(Gather), MixedUnionNode, null-aware right semi project
HashJoinNode, and TableWriteNode(!supportsMultiThreading).

Refactor LocalPlanner::maxDrivers to use the new API instead of per-type
dynamic_casts for single-thread checks. Config-dependent logic (writer
counts, repartition partition count) remains unchanged.

Differential Revision: D96366381

Author: mbasmanova

velox/core/PlanNode.h

@@ -206,6 +206,14 @@ class PlanNode : public ISerializable {
     return false;
   }
 
+  /// Returns true if this plan node requires single-threaded execution
+  /// (maxDrivers = 1). For example, ValuesNode, final OrderByNode, final
+  /// LimitNode, MergeExchangeNode, LocalMergeNode, and
+  /// LocalPartitionNode(Gather) all require single-threaded execution.
+  virtual bool requiresSingleThread() const {
+    return false;
+  }
+
   /// Returns true if this plan node operator supports task barrier processing.
   /// To support barrier processing, the operator must be able to drain its
   /// buffered output when it receives the drain signal at split boundary. Not
@@ -391,6 +399,10 @@ class ValuesNode : public PlanNode {
   void accept(const PlanNodeVisitor& visitor, PlanNodeVisitorContext& context)
       const override;
 
+  bool requiresSingleThread() const override {
+    return !parallelizable_;
+  }
+
   const std::vector<RowVectorPtr>& values() const {
     return values_;
   }
@@ -493,6 +505,10 @@ class ArrowStreamNode : public PlanNode {
   void accept(const PlanNodeVisitor& visitor, PlanNodeVisitorContext& context)
       const override;
 
+  bool requiresSingleThread() const override {
+    return true;
+  }
+
   const std::shared_ptr<ArrowArrayStream>& arrowStream() const {
     return arrowStream_;
   }
@@ -1713,6 +1729,11 @@ class TableWriteNode : public PlanNode {
     return columnStatsSpec_;
   }
 
+  bool requiresSingleThread() const override {
+    return !insertTableHandle_->connectorInsertTableHandle()
+                ->supportsMultiThreading();
+  }
+
   bool canSpill(const QueryConfig& queryConfig) const override {
     return queryConfig.writerSpillEnabled();
   }
@@ -1846,6 +1867,10 @@ class TableWriteMergeNode : public PlanNode {
     return columnStatsSpec_;
   }
 
+  bool requiresSingleThread() const override {
+    return true;
+  }
+
   const std::vector<PlanNodePtr>& sources() const override {
     return sources_;
   }
@@ -2351,6 +2376,10 @@ class MergeExchangeNode : public ExchangeNode {
     return sortingOrders_;
   }
 
+  bool requiresSingleThread() const override {
+    return true;
+  }
+
   void accept(const PlanNodeVisitor& visitor, PlanNodeVisitorContext& context)
       const override;
 
@@ -2449,6 +2478,10 @@ class LocalMergeNode : public PlanNode {
   void accept(const PlanNodeVisitor& visitor, PlanNodeVisitorContext& context)
       const override;
 
+  bool requiresSingleThread() const override {
+    return true;
+  }
+
   const std::vector<FieldAccessTypedExprPtr>& sortingKeys() const {
     return sortingKeys_;
   }
@@ -2660,6 +2693,10 @@ class LocalPartitionNode : public PlanNode {
     return scaleWriter_;
   }
 
+  bool requiresSingleThread() const override {
+    return type_ == Type::kGather;
+  }
+
   bool supportsBarrier() const override {
     return !scaleWriter_;
   }
@@ -3433,6 +3470,10 @@ class HashJoinNode : public AbstractJoinNode {
         queryConfig.joinSpillEnabled();
   }
 
+  bool requiresSingleThread() const override {
+    return isRightSemiProjectJoin() && nullAware_;
+  }
+
   bool isNullAware() const {
     return nullAware_;
   }
@@ -3507,6 +3548,10 @@ class MergeJoinNode : public AbstractJoinNode {
     }
   };
 
+  bool requiresSingleThread() const override {
+    return true;
+  }
+
   std::string_view name() const override {
     return "MergeJoin";
   }
@@ -4084,6 +4129,10 @@ class OrderByNode : public PlanNode {
   void accept(const PlanNodeVisitor& visitor, PlanNodeVisitorContext& context)
       const override;
 
+  bool requiresSingleThread() const override {
+    return !isPartial_;
+  }
+
   // True if this node only sorts a portion of the final result. If it is
   // true, a local merge or merge exchange is required to merge the sorted
   // runs.
@@ -4431,6 +4480,10 @@ class TopNNode : public PlanNode {
   void accept(const PlanNodeVisitor& visitor, PlanNodeVisitorContext& context)
       const override;
 
+  bool requiresSingleThread() const override {
+    return !isPartial_;
+  }
+
   int32_t count() const {
     return count_;
   }
@@ -4566,6 +4619,10 @@ class LimitNode : public PlanNode {
     return count_;
   }
 
+  bool requiresSingleThread() const override {
+    return !isPartial_;
+  }
+
   bool isPartial() const {
     return isPartial_;
   }
@@ -5843,6 +5900,10 @@ class MixedUnionNode : public PlanNode {
     return batchSizesPerSource_[sourceIndex];
   }
 
+  bool requiresSingleThread() const override {
+    return true;
+  }
+
   void accept(const PlanNodeVisitor& visitor, PlanNodeVisitorContext& context)
       const override;
 

velox/exec/LocalPlanner.cpp

@@ -283,8 +283,7 @@ void plan(
 // Sometimes consumer limits the number of drivers its producer can run.
 uint32_t maxDriversForConsumer(
     const std::shared_ptr<const core::PlanNode>& node) {
-  if (std::dynamic_pointer_cast<const core::MergeJoinNode>(node)) {
-    // MergeJoinNode must run single-threaded.
+  if (node && node->requiresSingleThread()) {
     return 1;
   }
   return std::numeric_limits<uint32_t>::max();
@@ -298,83 +297,24 @@ uint32_t maxDrivers(
     return count;
   }
   for (auto& node : driverFactory.planNodes) {
-    if (auto topN = std::dynamic_pointer_cast<const core::TopNNode>(node)) {
-      if (!topN->isPartial()) {
-        // final topN must run single-threaded
-        return 1;
-      }
-    } else if (
-        auto values = std::dynamic_pointer_cast<const core::ValuesNode>(node)) {
-      // values node must run single-threaded, unless in test context
-      if (!values->testingIsParallelizable()) {
-        return 1;
-      }
-    } else if (std::dynamic_pointer_cast<const core::ArrowStreamNode>(node)) {
-      // ArrowStream node must run single-threaded.
+    if (node->requiresSingleThread()) {
       return 1;
-    } else if (
-        auto limit = std::dynamic_pointer_cast<const core::LimitNode>(node)) {
-      // final limit must run single-threaded
-      if (!limit->isPartial()) {
-        return 1;
-      }
-    } else if (
-        auto orderBy =
-            std::dynamic_pointer_cast<const core::OrderByNode>(node)) {
-      // final orderby must run single-threaded
-      if (!orderBy->isPartial()) {
-        return 1;
-      }
-    } else if (
-        auto localExchange =
+    }
+
+    if (auto localExchange =
             std::dynamic_pointer_cast<const core::LocalPartitionNode>(node)) {
-      // Local gather must run single-threaded.
-      switch (localExchange->type()) {
-        case core::LocalPartitionNode::Type::kGather:
-          return 1;
-        case core::LocalPartitionNode::Type::kRepartition:
-          count = std::min(queryConfig.maxLocalExchangePartitionCount(), count);
-          break;
-        default:
-          VELOX_UNREACHABLE("Unexpected local exchange type");
+      // Repartition limits parallelism to the partition count.
+      if (localExchange->type() ==
+          core::LocalPartitionNode::Type::kRepartition) {
+        count = std::min(queryConfig.maxLocalExchangePartitionCount(), count);
       }
-    } else if (std::dynamic_pointer_cast<const core::LocalMergeNode>(node)) {
-      // Local merge must run single-threaded.
-      return 1;
-    } else if (std::dynamic_pointer_cast<const core::MixedUnionNode>(node)) {
-      // Mixed union must run single-threaded.
-      return 1;
-    } else if (std::dynamic_pointer_cast<const core::MergeExchangeNode>(node)) {
-      // Merge exchange must run single-threaded.
-      return 1;
-    } else if (std::dynamic_pointer_cast<const core::MergeJoinNode>(node)) {
-      // Merge join must run single-threaded.
-      return 1;
-    } else if (
-        auto join = std::dynamic_pointer_cast<const core::HashJoinNode>(node)) {
-      // Null-aware right semi project doesn't support multi-threaded
-      // execution.
-      if (join->isRightSemiProjectJoin() && join->isNullAware()) {
-        return 1;
-      }
-    } else if (std::dynamic_pointer_cast<const core::TableWriteMergeNode>(
-                   node)) {
-      // TableWriteMerge accumulates state (row counts, fragments, stats)
-      // and produces a single summary row. Must run single-threaded.
-      return 1;
     } else if (
         auto tableWrite =
             std::dynamic_pointer_cast<const core::TableWriteNode>(node)) {
-      const auto& connectorInsertHandle =
-          tableWrite->insertTableHandle()->connectorInsertTableHandle();
-      if (!connectorInsertHandle->supportsMultiThreading()) {
-        return 1;
+      if (tableWrite->hasPartitioningScheme()) {
+        return queryConfig.taskPartitionedWriterCount();
       } else {
-        if (tableWrite->hasPartitioningScheme()) {
-          return queryConfig.taskPartitionedWriterCount();
-        } else {
-          return queryConfig.taskWriterCount();
-        }
+        return queryConfig.taskWriterCount();
       }
     } else {
       auto result = Operator::maxDrivers(node);