Optimized PartitionedOutput staging hub

scripts/setup-common.sh

@@ -48,7 +48,12 @@ function install_fmt {
 
 function install_folly {
   wget_and_untar https://github.com/facebook/folly/archive/refs/tags/"${FB_OS_VERSION}".tar.gz folly
-  local FOLLY_FLAGS=(-DBUILD_SHARED_LIBS="$VELOX_BUILD_SHARED" -DBUILD_TESTS=OFF -DFOLLY_HAVE_INT128_T=ON)
+  local FOLLY_FLAGS=(
+    -DBUILD_SHARED_LIBS="$VELOX_BUILD_SHARED"
+    -DBUILD_TESTS=OFF
+    -DCMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE}"
+    -DFOLLY_HAVE_INT128_T=ON
+  )
   # When folly is static, use static gflags to avoid dual gflags flag
   # registration when .so plugins are dlopen'd (both the binary and plugin
   # would register the same flags in a shared gflags registry).

scripts/setup-helper-functions.sh

@@ -81,7 +81,8 @@ function github_checkout {
 # The values that CPU_ARCH can take are as follows:
 #   arm64  : Target Apple silicon.
 #   aarch64: Target general 64 bit arm cpus.
-#   avx:     Target Intel CPUs with AVX.
+#   avx512:  Target Intel CPUs with AVX-512F.
+#   avx:     Target Intel CPUs with AVX2.
 #   sse:     Target Intel CPUs with sse.
 # Echo's the appropriate compiler flags which can be captured as so
 # CXX_FLAGS=$(get_cxx_flags) or
@@ -102,7 +103,9 @@ function get_cxx_flags {
       else # x86_64
         local CPU_CAPABILITIES
         CPU_CAPABILITIES=$(sysctl -a | grep machdep.cpu.features | awk '{print tolower($0)}')
-        if [[ $CPU_CAPABILITIES =~ "avx" ]]; then
+        if [[ $CPU_CAPABILITIES =~ "avx512f" ]]; then
+          CPU_ARCH="avx512"
+        elif [[ $CPU_CAPABILITIES =~ "avx" ]]; then
           CPU_ARCH="avx"
         else
           CPU_ARCH="sse"
@@ -114,7 +117,9 @@ function get_cxx_flags {
       else # x86_64
         local CPU_CAPABILITIES
         CPU_CAPABILITIES=$(cat /proc/cpuinfo | grep flags | head -n 1 | awk '{print tolower($0)}')
-        if [[ $CPU_CAPABILITIES =~ "avx" ]]; then
+        if [[ $CPU_CAPABILITIES =~ "avx512f" ]]; then
+          CPU_ARCH="avx512"
+        elif [[ $CPU_CAPABILITIES =~ "avx" ]]; then
           CPU_ARCH="avx"
         elif [[ $CPU_CAPABILITIES =~ "sse" ]]; then
           CPU_ARCH="sse"
@@ -131,8 +136,12 @@ function get_cxx_flags {
     echo -n "-mcpu=apple-m1+crc"
     ;;
 
+  "avx512")
+    echo -n "-mavx512f -mavx2 -mfma -mavx -mf16c -mlzcnt -mbmi2"
+    ;;
+
   "avx")
-    echo -n "-mavx2 -mfma -mavx -mf16c -mlzcnt  -mbmi2"
+    echo -n "-mavx2 -mfma -mavx -mf16c -mlzcnt -mbmi2"
     ;;
 
   "sse")

velox/common/process/ProcessBase.cpp

@@ -32,6 +32,8 @@ DECLARE_bool(avx2); // Enables use of AVX2 when available NOLINT
 
 DECLARE_bool(bmi2); // Enables use of BMI2 when available NOLINT
 
+DECLARE_bool(avx512f);
+
 namespace facebook {
 namespace velox {
 namespace process {
@@ -106,6 +108,7 @@ uint64_t threadCpuNanos() {
 namespace {
 bool bmi2CpuFlag = folly::CpuId().bmi2();
 bool avx2CpuFlag = folly::CpuId().avx2();
+bool avx512fCpuFlag = folly::CpuId().avx512f();
 } // namespace
 
 bool hasAvx2() {
@@ -124,6 +127,14 @@ bool hasBmi2() {
 #endif
 }
 
+bool hasAvx512f() {
+#ifdef __AVX512F__
+  return avx512fCpuFlag && FLAGS_avx512f;
+#else
+  return false;
+#endif
+}
+
 } // namespace process
 } // namespace velox
 } // namespace facebook

velox/common/process/ProcessBase.h

@@ -46,6 +46,10 @@ uint64_t threadCpuNanos();
 /// by flag.
 bool hasAvx2();
 
+/// True if the machine has Intel AVX512F instructions and these are not
+/// disabled by flag.
+bool hasAvx512f();
+
 /// True if the machine has Intel BMI2 instructions and these are not disabled
 /// by flag.
 bool hasBmi2();

velox/connectors/hive/HiveConnector.cpp

@@ -132,7 +132,8 @@ void HiveConnector::registerSerDe() {
 
 std::unique_ptr<core::PartitionFunction> HivePartitionFunctionSpec::create(
     int numPartitions,
-    bool localExchange) const {
+    bool localExchange,
+    bool /*useOptimizedPartitionFunction*/) const {
   std::vector<int> bucketToPartitions;
   if (bucketToPartition_.empty()) {
     // NOTE: if hive partition function spec doesn't specify bucket to partition

velox/connectors/hive/HiveConnector.h

@@ -141,7 +141,8 @@ class HivePartitionFunctionSpec : public core::PartitionFunctionSpec {
 
   std::unique_ptr<core::PartitionFunction> create(
       int numPartitions,
-      bool localExchange) const override;
+      bool localExchange,
+      bool useOptimizedPartitionFunction = false) const override;
 
   std::string toString() const override;
 

velox/core/PlanNode.h

@@ -2500,9 +2500,13 @@ class PartitionFunctionSpec : public ISerializable {
  public:
   /// If 'localExchange' is true, the partition function is used for local
   /// exchange within a velox task.
+  /// TODO: useOptimizedPartitionFunction = true is only supported in
+  /// HashPartitionFunction now. Will extend the optimization to other
+  /// PartitionFunctions soon.
   virtual std::unique_ptr<PartitionFunction> create(
       int numPartitions,
-      bool localExchange = false) const = 0;
+      bool localExchange = false,
+      bool useOptimizedPartitionFunction = false) const = 0;
 
   virtual ~PartitionFunctionSpec() = default;
 
@@ -2515,7 +2519,8 @@ class GatherPartitionFunctionSpec : public PartitionFunctionSpec {
  public:
   std::unique_ptr<PartitionFunction> create(
       int /*numPartitions*/,
-      bool /*localExchange*/) const override {
+      bool /*localExchange*/,
+      bool /*useOptimizedPartitionFunction*/ = false) const override {
     VELOX_UNREACHABLE();
   }
 

velox/core/QueryConfig.cpp

@@ -90,6 +90,7 @@ const std::vector<config::ConfigProperty>& QueryConfig::registeredProperties() {
 
     // Partitioned output.
     VELOX_REGISTER_QUERY_CONFIG(kPartitionedOutputEagerFlush);
+    VELOX_REGISTER_QUERY_CONFIG(kOptimizedHashPartitionFunctionEnabled);
     VELOX_REGISTER_QUERY_CONFIG(kMaxPartitionedOutputBufferSize);
     VELOX_REGISTER_QUERY_CONFIG(kMaxOutputBufferSize);
 

velox/core/QueryConfig.h

@@ -454,6 +454,16 @@ class QueryConfig {
       false,
       "Flush PartitionedOutput rows eagerly without buffering.")
 
+  /// If true, use OptimizedHashPartitionFunction in place of
+  /// HashPartitionFunction.
+  VELOX_QUERY_CONFIG(
+      kOptimizedHashPartitionFunctionEnabled,
+      optimizedHashPartitionFunctionEnabled,
+      "optimized_hash_partition_function_enabled",
+      bool,
+      false,
+      "Use OptimizedHashPartitionFunction instead of HashPartitionFunction.")
+
   /// The maximum number of bytes to buffer in PartitionedOutput operator to
   /// avoid creating tiny SerializedPages.
   VELOX_QUERY_CONFIG(
@@ -1469,6 +1479,14 @@ class QueryConfig {
       1000,
       "Batch size threshold for zero-copy in MarkSorted operator.")
 
+  VELOX_QUERY_CONFIG(
+      kOptimizedPartitionedOutputEnabled,
+      optimizedPartitionedOutputEnabled,
+      "optimized_repartitioning",
+      bool,
+      false,
+      "Enable OptimizedPartitionedOutput operator.");
+
   // --- Hand-written accessors for properties that need custom logic ---
 
   // Generated by VELOX_QUERY_CONFIG for simple properties above.

velox/exec/CMakeLists.txt

@@ -71,6 +71,9 @@ velox_add_library(
   OperatorTraceScan.cpp
   OperatorTraceWriter.cpp
   OperatorUtils.cpp
+  OptimizedHashPartitionFunction.cpp
+  OptimizedPartitionedOutput.cpp
+  OptimizedVectorHasher.cpp
   OrderBy.cpp
   OutputBuffer.cpp
   OutputBufferManager.cpp
@@ -177,6 +180,7 @@ velox_add_library(
   OperatorTraceWriter.h
   OperatorType.h
   OperatorUtils.h
+  OptimizedVectorHasher.h
   OrderBy.h
   OutputBuffer.h
   OutputBufferManager.h

velox/exec/HashPartitionFunction.cpp

@@ -13,8 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#include <velox/exec/HashPartitionFunction.h>
-#include <velox/exec/VectorHasher.h>
+#include "velox/exec/HashPartitionFunction.h"
+
+#include "velox/exec/OptimizedHashPartitionFunction.h"
+#include "velox/exec/VectorHasher.h"
 
 #define XXH_INLINE_ALL
 #include <xxhash.h> // @manual=third-party//xxHash:xxhash
@@ -123,9 +125,15 @@ std::optional<uint32_t> HashPartitionFunction::partition(
 
 std::unique_ptr<core::PartitionFunction> HashPartitionFunctionSpec::create(
     int numPartitions,
-    bool localExchange) const {
-  return std::make_unique<exec::HashPartitionFunction>(
-      localExchange, numPartitions, inputType_, keyChannels_, constValues_);
+    bool localExchange,
+    bool useOptimizedPartitionFunction) const {
+  return createHashPartitionFunction(
+      localExchange,
+      numPartitions,
+      inputType_,
+      keyChannels_,
+      constValues_,
+      useOptimizedPartitionFunction);
 }
 
 std::string HashPartitionFunctionSpec::toString() const {
@@ -180,4 +188,33 @@ core::PartitionFunctionSpecPtr HashPartitionFunctionSpec::deserialize(
   return std::make_shared<HashPartitionFunctionSpec>(
       ISerializable::deserialize<RowType>(obj["inputType"]), keys, constValues);
 }
+
+std::unique_ptr<HashPartitionFunctionBase> createHashPartitionFunction(
+    bool localExchange,
+    int numPartitions,
+    const RowTypePtr& inputType,
+    const std::vector<column_index_t>& keyChannels,
+    const std::vector<VectorPtr>& constValues,
+    bool useOptimizedPartitionFunction) {
+  if (useOptimizedPartitionFunction) {
+    return std::make_unique<OptimizedHashPartitionFunction>(
+        localExchange, numPartitions, inputType, keyChannels, constValues);
+  }
+  return std::make_unique<HashPartitionFunction>(
+      localExchange, numPartitions, inputType, keyChannels, constValues);
+}
+
+std::unique_ptr<HashPartitionFunctionBase> createHashPartitionFunction(
+    const HashBitRange& hashBitRange,
+    const RowTypePtr& inputType,
+    const std::vector<column_index_t>& keyChannels,
+    const std::vector<VectorPtr>& constValues,
+    bool useOptimizedPartitionFunction) {
+  if (useOptimizedPartitionFunction) {
+    return std::make_unique<OptimizedHashPartitionFunction>(
+        hashBitRange, inputType, keyChannels, constValues);
+  }
+  return std::make_unique<HashPartitionFunction>(
+      hashBitRange, inputType, keyChannels, constValues);
+}
 } // namespace facebook::velox::exec

velox/exec/HashPartitionFunction.h

@@ -15,19 +15,28 @@
  */
 #pragma once
 
-#include <velox/exec/HashBitRange.h>
-#include <velox/exec/VectorHasher.h>
 #include "velox/core/PlanNode.h"
+#include "velox/exec/HashBitRange.h"
+#include "velox/exec/VectorHasher.h"
 
 namespace facebook::velox::exec {
 
+class HashPartitionFunctionBase : public core::PartitionFunction {
+ public:
+  ~HashPartitionFunctionBase() override = default;
+
+  virtual int numPartitions() const = 0;
+};
+
 /// Calculates partition number for each row of the specified vector using a
 /// hash function. The constructor with hashBitRange parameter requires both
 /// hashBitRange and keyChannels to be non-empty. The constructor with
 /// numPartitions allows the keyChannels argument to be empty. If keyChannels is
 /// empty, then the resulting partition number of partition() will always be
 /// zero.
-class HashPartitionFunction : public core::PartitionFunction {
+/// Extends PartitionFunction with access to the configured number of
+/// partitions.
+class HashPartitionFunction : public HashPartitionFunctionBase {
  public:
   HashPartitionFunction(
       bool localExchange,
@@ -48,7 +57,7 @@ class HashPartitionFunction : public core::PartitionFunction {
       const RowVector& input,
       std::vector<uint32_t>& partitions) override;
 
-  int numPartitions() const {
+  int numPartitions() const override {
     return numPartitions_;
   }
 
@@ -85,7 +94,8 @@ class HashPartitionFunctionSpec : public core::PartitionFunctionSpec {
 
   std::unique_ptr<core::PartitionFunction> create(
       int numPartitions,
-      bool localExchange) const override;
+      bool localExchange,
+      bool useOptimizedPartitionFunction = false) const override;
 
   std::string toString() const override;
 
@@ -100,4 +110,22 @@ class HashPartitionFunctionSpec : public core::PartitionFunctionSpec {
   const std::vector<column_index_t> keyChannels_;
   const std::vector<VectorPtr> constValues_;
 };
+
+/// Creates either HashPartitionFunction or OptimizedHashPartitionFunction
+/// based on 'useOptimizedPartitionFunction'.
+std::unique_ptr<HashPartitionFunctionBase> createHashPartitionFunction(
+    bool localExchange,
+    int numPartitions,
+    const RowTypePtr& inputType,
+    const std::vector<column_index_t>& keyChannels,
+    const std::vector<VectorPtr>& constValues = {},
+    bool useOptimizedPartitionFunction = false);
+
+std::unique_ptr<HashPartitionFunctionBase> createHashPartitionFunction(
+    const HashBitRange& hashBitRange,
+    const RowTypePtr& inputType,
+    const std::vector<column_index_t>& keyChannels,
+    const std::vector<VectorPtr>& constValues = {},
+    bool useOptimizedPartitionFunction = false);
+
 } // namespace facebook::velox::exec

velox/exec/LocalPartition.cpp

@@ -339,10 +339,13 @@ LocalPartition::LocalPartition(
           ctx->task->getLocalExchangeQueues(ctx->splitGroupId, planNode->id())},
       numPartitions_{queues_.size()},
       partitionFunction_(
-          numPartitions_ == 1 ? nullptr
-                              : planNode->partitionFunctionSpec().create(
-                                    numPartitions_,
-                                    /*localExchange=*/true)),
+          numPartitions_ == 1
+              ? nullptr
+              : planNode->partitionFunctionSpec().create(
+                    numPartitions_,
+                    /*localExchange=*/true,
+                    ctx->queryConfig()
+                        .optimizedHashPartitionFunctionEnabled())),
       singlePartitionBufferSize_{
           (numPartitions_ <
                ctx->queryConfig()

velox/exec/LocalPlanner.cpp

@@ -37,6 +37,7 @@
 #include "velox/exec/NestedLoopJoinBuild.h"
 #include "velox/exec/NestedLoopJoinProbe.h"
 #include "velox/exec/OperatorTraceScan.h"
+#include "velox/exec/OptimizedPartitionedOutput.h"
 #include "velox/exec/OrderBy.h"
 #include "velox/exec/ParallelProject.h"
 #include "velox/exec/PartitionedOutput.h"
@@ -553,9 +554,15 @@ std::shared_ptr<Driver> DriverFactory::createDriver(
         auto partitionedOutputNode =
             std::dynamic_pointer_cast<const core::PartitionedOutputNode>(
                 planNode)) {
-      operators.push_back(
-          std::make_unique<PartitionedOutput>(
-              id, ctx.get(), partitionedOutputNode, eagerFlush(*planNode)));
+      if (ctx->queryConfig().optimizedPartitionedOutputEnabled()) {
+        operators.push_back(
+            std::make_unique<OptimizedPartitionedOutput>(
+                id, ctx.get(), partitionedOutputNode));
+      } else {
+        operators.push_back(
+            std::make_unique<PartitionedOutput>(
+                id, ctx.get(), partitionedOutputNode, eagerFlush(*planNode)));
+      }
     } else if (
         auto joinNode =
             std::dynamic_pointer_cast<const core::HashJoinNode>(planNode)) {