feat(cudf): Add round robin batch support for Local partition in Velox-cudf (facebookincubator#14664)

Summary:
Round Robin batch partitioning can be supported with CPU LocalPartition operator itself. This PR adds changes to use LocalPartition only for Hash, and use LocalPartition for RoundRobin. Unit tests are added to test RoundRobin under difference configurations.
Unit tests
- roundRobinMultipleBatches
- roundRobinEmptyInput
- roundRobinMultipleSources
- roundRobinWithAggregation
- roundRobinWithTableScan
- roundRobinAllCombinations

This PR also has related fixes to CudfPlanBuilder after commit facebookincubator@31200dd which affects 3 function interfaces.

Pull Request resolved: facebookincubator#14664

Reviewed By: Yuhta

Differential Revision: D83667533

Pulled By: kgpai

fbshipit-source-id: c8e555f81ee3a8fcd69d5f47c34d140238f894d3

Author: karthikeyann

velox/experimental/cudf/exec/CudfLocalPartition.cpp

@@ -17,13 +17,23 @@
 #include "velox/experimental/cudf/exec/CudfLocalPartition.h"
 #include "velox/experimental/cudf/vector/CudfVector.h"
 
+#include "velox/exec/HashPartitionFunction.h"
 #include "velox/exec/Task.h"
 
 #include <cudf/copying.hpp>
 #include <cudf/partitioning.hpp>
 
 namespace facebook::velox::cudf_velox {
 
+bool CudfLocalPartition::shouldReplace(
+    const std::shared_ptr<const core::LocalPartitionNode>& planNode) {
+  auto* hashFunctionSpec = dynamic_cast<const exec::HashPartitionFunctionSpec*>(
+      &planNode->partitionFunctionSpec());
+  // Only replace LocalPartition with CudfLocalPartition for hash partitioning.
+  // TODO: Round Robin Row-Wise Partitioning can be supported in future.
+  return hashFunctionSpec;
+}
+
 CudfLocalPartition::CudfLocalPartition(
     int32_t operatorId,
     exec::DriverCtx* ctx,
@@ -55,9 +65,11 @@ CudfLocalPartition::CudfLocalPartition(
 
   // Get partition function specification string
   std::string spec = planNode->partitionFunctionSpec().toString();
+  auto* hashFunctionSpec = dynamic_cast<const exec::HashPartitionFunctionSpec*>(
+      &planNode->partitionFunctionSpec());
 
   // Only parse keys if it's a hash function
-  if (spec.find("HASH(") != std::string::npos) {
+  if (hashFunctionSpec) {
     // Extract keys between HASH( and )
     size_t start = spec.find("HASH(") + 5;
     size_t end = spec.find(")", start);

velox/experimental/cudf/exec/CudfLocalPartition.h

@@ -51,6 +51,9 @@ class CudfLocalPartition : public exec::Operator, public NvtxHelper {
 
   bool isFinished() override;
 
+  static bool shouldReplace(
+      const std::shared_ptr<const core::LocalPartitionNode>& planNode);
+
  protected:
   const std::vector<std::shared_ptr<exec::LocalExchangeQueue>> queues_;
   const size_t numPartitions_;

velox/experimental/cudf/exec/ToCudf.cpp

@@ -206,6 +206,18 @@ bool CompileState::compile(bool force_replace) {
           id, planNode->outputType(), ctx, planNode->id() + "-from-velox"));
       replaceOp.back()->initialize();
     }
+    if (not replaceOp.empty()) {
+      // from-velox only, because need to inserted before current operator.
+      operatorsOffset += replaceOp.size();
+      [[maybe_unused]] auto replaced = driverFactory_.replaceOperators(
+          driver_,
+          replacingOperatorIndex,
+          replacingOperatorIndex,
+          std::move(replaceOp));
+      replacingOperatorIndex = operatorIndex + operatorsOffset;
+      replaceOp.clear();
+      replacementsMade = true;
+    }
 
     // This is used to denote if the current operator is kept or replaced.
     auto keepOperator = 0;
@@ -273,9 +285,15 @@ bool CompileState::compile(bool force_replace) {
       auto planNode = std::dynamic_pointer_cast<const core::LocalPartitionNode>(
           getPlanNode(localPartitionOp->planNodeId()));
       VELOX_CHECK(planNode != nullptr);
-      replaceOp.push_back(
-          std::make_unique<CudfLocalPartition>(id, ctx, planNode));
-      replaceOp.back()->initialize();
+      if (CudfLocalPartition::shouldReplace(planNode)) {
+        replaceOp.push_back(
+            std::make_unique<CudfLocalPartition>(id, ctx, planNode));
+        replaceOp.back()->initialize();
+      } else {
+        // Round Robin batch-wise Partitioning is supported by CPU operator with
+        // GPU Vector.
+        keepOperator = 1;
+      }
     } else if (
         auto localExchangeOp = dynamic_cast<exec::LocalExchange*>(oper)) {
       keepOperator = 1;
@@ -319,6 +337,7 @@ bool CompileState::compile(bool force_replace) {
     }
 
     if (not replaceOp.empty()) {
+      // ReplaceOp, to-velox.
       operatorsOffset += replaceOp.size() - 1 + keepOperator;
       [[maybe_unused]] auto replaced = driverFactory_.replaceOperators(
           driver_,