From 37e5ae94a7a4717dcdcbad40bb9629b9f1f5fbdc Mon Sep 17 00:00:00 2001
From: Andrii Rosa <andriirosa@meta.com>
Date: Fri, 6 Dec 2024 06:39:22 -0800
Subject: [PATCH] feat: Allow non standard partition functions in
 ScaleWriterPartitioningLocalPartition (#11762)

Summary:

For example Hive connector partitioning function

Reviewed By: xiaoxmeng

Differential Revision: D66833831
---
 velox/exec/LocalPlanner.cpp              | 8 +++-----
 velox/exec/ScaleWriterLocalPartition.cpp | 4 ----
 2 files changed, 3 insertions(+), 9 deletions(-)

diff --git a/velox/exec/LocalPlanner.cpp b/velox/exec/LocalPlanner.cpp
index 3645e737b83e8..8112318819235 100644
--- a/velox/exec/LocalPlanner.cpp
+++ b/velox/exec/LocalPlanner.cpp
@@ -76,15 +76,13 @@ std::unique_ptr<Operator> createScaleWriterLocalPartition(
     const std::shared_ptr<const core::LocalPartitionNode>& localPartitionNode,
     int32_t operatorId,
     DriverCtx* ctx) {
-  if (dynamic_cast<const HashPartitionFunctionSpec*>(
+  if (dynamic_cast<const RoundRobinPartitionFunctionSpec*>(
           &localPartitionNode->partitionFunctionSpec())) {
-    return std::make_unique<ScaleWriterPartitioningLocalPartition>(
+    return std::make_unique<ScaleWriterLocalPartition>(
         operatorId, ctx, localPartitionNode);
   }
 
-  VELOX_CHECK_NOT_NULL(dynamic_cast<const RoundRobinPartitionFunctionSpec*>(
-      &localPartitionNode->partitionFunctionSpec()));
-  return std::make_unique<ScaleWriterLocalPartition>(
+  return std::make_unique<ScaleWriterPartitioningLocalPartition>(
       operatorId, ctx, localPartitionNode);
 }
 
diff --git a/velox/exec/ScaleWriterLocalPartition.cpp b/velox/exec/ScaleWriterLocalPartition.cpp
index c4995d0975281..b243e2b808bb6 100644
--- a/velox/exec/ScaleWriterLocalPartition.cpp
+++ b/velox/exec/ScaleWriterLocalPartition.cpp
@@ -57,10 +57,6 @@ ScaleWriterPartitioningLocalPartition::ScaleWriterPartitioningLocalPartition(
       : planNode->partitionFunctionSpec().create(
             numTablePartitions_,
             /*localExchange=*/true);
-  if (partitionFunction_ != nullptr) {
-    VELOX_CHECK_NOT_NULL(
-        dynamic_cast<HashPartitionFunction*>(partitionFunction_.get()));
-  }
 }
 
 void ScaleWriterPartitioningLocalPartition::initialize() {