feat: Add SpillPartitionFunction (facebookincubator#13079)

Summary:
Add spill partition function later used for mixed mode hash join spilling
Pull Request resolved: facebookincubator#13079

Reviewed By: xiaoxmeng

Differential Revision: D73339693

fbshipit-source-id: 857f4e9dbb4f11b46e3154131d9dc5c62611963d

Author: tanjialiang

velox/exec/Spill.cpp

@@ -523,6 +523,41 @@ SpillPartitionId SpillPartitionIdLookup::partition(uint64_t hash) const {
   return lookup_[bits::extractBits(hash, partitionBitsMask_)];
 }
 
+SpillPartitionFunction::SpillPartitionFunction(
+    const SpillPartitionIdLookup& lookup,
+    const RowTypePtr& inputType,
+    const std::vector<column_index_t>& keyChannels)
+    : lookup_(lookup) {
+  VELOX_CHECK(!keyChannels.empty(), "Key channels must not be empty.");
+  hashers_.reserve(keyChannels.size());
+  for (const auto channel : keyChannels) {
+    VELOX_CHECK_NE(channel, kConstantChannel);
+    hashers_.emplace_back(
+        VectorHasher::create(inputType->childAt(channel), channel));
+  }
+}
+
+void SpillPartitionFunction::partition(
+    const RowVector& input,
+    std::vector<SpillPartitionId>& partitionIds) {
+  const auto size = input.size();
+  rows_.resize(size);
+  rows_.setAll();
+
+  hashes_.resize(size);
+  for (auto i = 0; i < hashers_.size(); ++i) {
+    auto& hasher = hashers_[i];
+    hashers_[i]->decode(*input.childAt(hasher->channel()), rows_);
+    hashers_[i]->hash(rows_, i > 0, hashes_);
+  }
+
+  partitionIds.resize(size);
+
+  for (auto i = 0; i < size; ++i) {
+    partitionIds[i] = lookup_.partition(hashes_[i]);
+  }
+}
+
 uint8_t partitionBitOffset(
     const SpillPartitionId& id,
     uint8_t startPartitionBitOffset,

velox/exec/Spill.h

@@ -27,11 +27,14 @@
 #include "velox/exec/SpillFile.h"
 #include "velox/exec/TreeOfLosers.h"
 #include "velox/exec/UnorderedStreamReader.h"
+#include "velox/exec/VectorHasher.h"
 #include "velox/vector/ComplexVector.h"
 #include "velox/vector/DecodedVector.h"
 #include "velox/vector/VectorStream.h"
 
 namespace facebook::velox::exec {
+class VectorHasher;
+
 /// A source of sorted spilled RowVectors coming either from a file or memory.
 class SpillMergeStream : public MergeStream {
  public:
@@ -212,7 +215,7 @@ class SpillPartitionId {
   /// 'SpillPartitionId'.
   static constexpr uint32_t kMaxPartitionBits{3};
 
-  /// Constructs an default invalid id.
+  /// Constructs a default invalid id.
   SpillPartitionId() = default;
 
   /// Constructs a root spill level id.
@@ -323,6 +326,30 @@ class SpillPartitionIdLookup {
   std::vector<SpillPartitionId> lookup_;
 };
 
+/// Vectorized partitioning function for spill. The partitioning takes advantage
+/// of SpillPartitionIdLookup and performs a fast partitioning for the input
+/// vector.
+class SpillPartitionFunction {
+ public:
+  SpillPartitionFunction(
+      const SpillPartitionIdLookup& lookup,
+      const RowTypePtr& inputType,
+      const std::vector<column_index_t>& keyChannels);
+
+  void partition(
+      const RowVector& input,
+      std::vector<SpillPartitionId>& partitionIds);
+
+ private:
+  const SpillPartitionIdLookup lookup_;
+
+  std::vector<std::unique_ptr<VectorHasher>> hashers_;
+
+  // Reusable resources for hashing.
+  SelectivityVector rows_;
+  raw_vector<uint64_t> hashes_;
+};
+
 /// Contains a spill partition data which includes the partition id and
 /// corresponding spill files.
 class SpillPartition {

velox/exec/tests/SpillTest.cpp

@@ -895,6 +895,118 @@ TEST_P(SpillTest, spillPartitionIdLookup) {
   }
 }
 
+TEST_P(SpillTest, spillPartitionFunctionBasic) {
+  std::vector<RowVectorPtr> inputVectors;
+  const auto rowType =
+      ROW({"key1", "key2", "key3", "value"},
+          {BIGINT(), VARCHAR(), VARCHAR(), VARCHAR()});
+  {
+    // Simple input vector for basic testing.
+    const uint64_t numRows = 100;
+    std::vector<VectorPtr> columns;
+    columns.push_back(
+        makeFlatVector<int64_t>(numRows, [](auto row) { return row; }));
+    columns.push_back(makeFlatVector<std::string>(
+        numRows, [](auto row) { return fmt::format("key_{}", row); }));
+    columns.push_back(makeFlatVector<std::string>(
+        numRows, [](auto row) { return fmt::format("key_{}_{}", row, row); }));
+    columns.push_back(makeFlatVector<std::string>(
+        numRows, [](auto row) { return fmt::format("val_{}", row); }));
+    inputVectors.push_back(makeRowVector(columns));
+  }
+
+  // Additional 2 fuzzed vector for wider range testing.
+  {
+    VectorFuzzer::Options options;
+    options.vectorSize = 100;
+    options.nullRatio = 0.3;
+    options.allowDictionaryVector = true;
+    inputVectors.push_back(VectorFuzzer(options, pool()).fuzzRow(rowType));
+  }
+
+  {
+    VectorFuzzer::Options options;
+    options.vectorSize = 1000;
+    options.nullRatio = 0.2;
+    options.stringVariableLength = true;
+    options.allowDictionaryVector = true;
+    inputVectors.push_back(VectorFuzzer(options, pool()).fuzzRow(rowType));
+  }
+
+  struct TestData {
+    std::string name;
+    SpillPartitionIdSet partitionIds;
+    uint32_t startPartitionBit;
+    uint32_t numPartitionBits;
+    std::vector<column_index_t> keyChannels;
+    size_t numRows;
+  };
+
+  std::vector<TestData> testCases = {
+      // Basic test with multiple partitions
+      {"basic",
+       {genPartitionId({0}),
+        genPartitionId({1}),
+        genPartitionId({2, 1}),
+        genPartitionId({3, 2})},
+       0,
+       2,
+       {0, 1},
+       100},
+
+      // Test with empty key channels
+      {"emptyKeys", {genPartitionId({0}), genPartitionId({1})}, 0, 2, {}, 50},
+
+      // Test with nulls
+      {"withNulls",
+       {genPartitionId({0}), genPartitionId({1}), genPartitionId({2})},
+       0,
+       2,
+       {0, 1},
+       80},
+
+      // Test with higher start partition bits
+      {"highStartBits",
+       {genPartitionId({0}),
+        genPartitionId({1}),
+        genPartitionId({2}),
+        genPartitionId({3})},
+       8,
+       2,
+       {0, 1, 2},
+       120}};
+
+  for (const auto& data : testCases) {
+    for (auto i = 0; i < inputVectors.size(); ++i) {
+      SCOPED_TRACE(fmt::format("Test case: {}, Input vector {}", data.name, i));
+
+      SpillPartitionIdLookup lookup(
+          data.partitionIds, data.startPartitionBit, data.numPartitionBits);
+      if (data.keyChannels.empty()) {
+        VELOX_ASSERT_THROW(
+            SpillPartitionFunction(lookup, rowType, data.keyChannels),
+            "Key channels must not be empty");
+        continue;
+      }
+      SpillPartitionFunction partitionFunction(
+          lookup, rowType, data.keyChannels);
+
+      auto& input = inputVectors[i];
+      std::vector<SpillPartitionId> resultPartitionIds;
+      partitionFunction.partition(*input, resultPartitionIds);
+
+      ASSERT_EQ(resultPartitionIds.size(), input->size());
+      for (const auto& id : resultPartitionIds) {
+        ASSERT_TRUE((!id.valid()) || data.partitionIds.contains(id));
+      }
+
+      std::unordered_set<SpillPartitionId> uniquePartitions(
+          resultPartitionIds.begin(), resultPartitionIds.end());
+      ASSERT_GT(uniquePartitions.size(), 1);
+    }
+  }
+}
+
 TEST_P(SpillTest, spillPartitionSet) {
   for (int iter = 0; iter < 3; ++iter) {
     folly::Random::DefaultGenerator rng;