Reused Hash Table in HashBuild Phase for Gluten

Author: JkSelf

velox/core/PlanNode.h

@@ -29,6 +29,12 @@ struct ArrowArrayStream;
 
 namespace facebook::velox::core {
 
+/// Direct access to HashTable from HashJoinNode would introduce a dependency
+/// cycle. We resolved this by defining an OpaqueHashTable interface,
+/// effectively decoupling the node logic from the specific table
+/// implementation.
+struct OpaqueHashTable;
+
 class PlanNodeVisitor;
 class PlanNodeVisitorContext;
 
@@ -3148,7 +3154,9 @@ class HashJoinNode : public AbstractJoinNode {
       PlanNodePtr left,
       PlanNodePtr right,
       RowTypePtr outputType,
-      bool useHashTableCache = false)
+      bool useHashTableCache = false,
+      bool joinHasNullKeys = false,
+      std::shared_ptr<OpaqueHashTable> reusableHashTable = nullptr)
       : AbstractJoinNode(
             id,
             joinType,
@@ -3159,7 +3167,9 @@ class HashJoinNode : public AbstractJoinNode {
             std::move(right),
             std::move(outputType)),
         nullAware_{nullAware},
-        useHashTableCache_{useHashTableCache} {
+        useHashTableCache_{useHashTableCache},
+        joinHasNullKeys_{joinHasNullKeys},
+        reusableHashTable_(std::move(reusableHashTable)) {
     validate();
 
     if (nullAware) {
@@ -3197,6 +3207,17 @@ class HashJoinNode : public AbstractJoinNode {
       return *this;
     }
 
+    Builder& joinHasNullKeys(bool joinHasNullKeys) {
+      joinHasNullKeys_ = joinHasNullKeys;
+      return *this;
+    }
+
+    Builder& reusableHashTable(
+        std::shared_ptr<OpaqueHashTable> opaqueHashTable) {
+      reusableHashTable_ = std::move(opaqueHashTable);
+      return *this;
+    }
+
     std::shared_ptr<HashJoinNode> build() const {
       VELOX_USER_CHECK(id_.has_value(), "HashJoinNode id is not set");
       VELOX_USER_CHECK(
@@ -3224,12 +3245,16 @@ class HashJoinNode : public AbstractJoinNode {
           left_.value(),
           right_.value(),
           outputType_.value(),
-          useHashTableCache_.value_or(false));
+          useHashTableCache_.value_or(false),
+          joinHasNullKeys_.value_or(false),
+          reusableHashTable_.value_or(nullptr));
     }
 
    private:
     std::optional<bool> nullAware_;
     std::optional<bool> useHashTableCache_;
+    std::optional<bool> joinHasNullKeys_;
+    std::optional<std::shared_ptr<OpaqueHashTable>> reusableHashTable_;
   };
 
   std::string_view name() const override {
@@ -3258,6 +3283,14 @@ class HashJoinNode : public AbstractJoinNode {
     return useHashTableCache_;
   }
 
+  bool joinHasNullKeys() const {
+    return joinHasNullKeys_;
+  }
+
+  std::shared_ptr<OpaqueHashTable> reusableHashTable() const {
+    return reusableHashTable_;
+  }
+
   folly::dynamic serialize() const override;
 
   static PlanNodePtr create(const folly::dynamic& obj, void* context);
@@ -3267,6 +3300,8 @@ class HashJoinNode : public AbstractJoinNode {
 
   const bool nullAware_;
   const bool useHashTableCache_;
+  const bool joinHasNullKeys_;
+  std::shared_ptr<OpaqueHashTable> reusableHashTable_;
 };
 
 using HashJoinNodePtr = std::shared_ptr<const HashJoinNode>;

velox/exec/HashBuild.cpp

@@ -47,6 +47,25 @@ BlockingReason fromStateToBlockingReason(HashBuild::State state) {
 }
 } // namespace
 
+void HashBuild::setReusableHashTable(
+    std::shared_ptr<core::OpaqueHashTable> opaqueHashTable) {
+  joinBridge_->start();
+  setState(State::kFinish);
+
+  if (joinNode_->joinHasNullKeys() && isAntiJoin(joinType_) && nullAware_ &&
+      !joinNode_->filter()) {
+    joinBridge_->setAntiJoinHasNullKeys();
+    return;
+  }
+
+  auto reusableHashTable =
+      std::reinterpret_pointer_cast<exec::BaseHashTable>(opaqueHashTable);
+
+  joinBridge_->setHashTable(
+      std::move(reusableHashTable), {}, joinNode_->joinHasNullKeys(), nullptr);
+  reuseHashTable_ = true;
+}
+
 HashBuild::HashBuild(
     int32_t operatorId,
     DriverCtx* driverCtx,
@@ -80,6 +99,12 @@ HashBuild::HashBuild(
 
   joinBridge_->addBuilder();
 
+  if (auto opaqueHashTable = joinNode_->reusableHashTable()) {
+    TestValue::adjust("facebook::velox::exec::HashBuild::HashBuild", this);
+    setReusableHashTable(opaqueHashTable);
+    return;
+  }
+
   const auto& inputType = joinNode_->sources()[1]->outputType();
 
   const auto numKeys = joinNode_->rightKeys().size();
@@ -121,7 +146,7 @@ HashBuild::HashBuild(
 void HashBuild::initialize() {
   Operator::initialize();
 
-  if (setupCachedHashTable()) {
+  if (setupCachedHashTable() || reuseHashTable_) {
     return;
   }
 

velox/exec/HashBuild.h

@@ -108,6 +108,10 @@ class HashBuild final : public Operator {
   // Invoked to set up hash table to build.
   void setupTable();
 
+  // Reuse the pre-built hash table.
+  void setReusableHashTable(
+      std::shared_ptr<core::OpaqueHashTable> opaqueHashTable);
+
   // Sets up hash table caching if enabled. Returns true if the cached table
   // is already available or if this operator should wait for another task
   // to build it, in which case further initialization should be skipped.
@@ -395,6 +399,8 @@ class HashBuild final : public Operator {
   // Count the number of hash table input rows for building deduped
   // hash table. It will not be updated after abandonBuildNoDupHash_ is true.
   int64_t numHashInputRows_ = 0;
+
+  bool reuseHashTable_ = false;
 };
 
 inline std::ostream& operator<<(std::ostream& os, HashBuild::State state) {

velox/exec/HashProbe.cpp

@@ -986,7 +986,8 @@ bool HashProbe::canSpill() const {
   // Hash table caching is incompatible with spilling. When the table is
   // cached and shared across tasks, clearing it after probe would corrupt
   // the cache for subsequent tasks.
-  if (joinNode_->useHashTableCache()) {
+  if (joinNode_->useHashTableCache() ||
+      joinNode_->reusableHashTable() != nullptr) {
     return false;
   }
   if (operatorCtx_->task()->hasMixedExecutionGroupJoin(joinNode_.get())) {

velox/exec/tests/HashJoinTest.cpp

@@ -8529,6 +8529,104 @@ TEST_P(HashJoinTest, arrayBasedLookupCustomComparisonType) {
   EXPECT_EQ(result->size(), 1'024);
 }
 
+DEBUG_ONLY_TEST_P(HashJoinTest, reuseHashTable) {
+  // Create build and probe vectors.
+  std::vector<RowVectorPtr> buildVectors = makeBatches(1, [&](int32_t) {
+    return makeRowVector(
+        {"u_0"},
+        {
+            makeFlatVector<int64_t>(100, [](auto row) { return row % 23; }),
+        });
+  });
+
+  std::vector<RowVectorPtr> probeVectors = makeBatches(5, [&](int32_t) {
+    return makeRowVector(
+        {"t_0"},
+        {
+            makeFlatVector<int64_t>(100, [](auto row) { return row % 23; }),
+        });
+  });
+
+  createDuckDbTable("t", probeVectors);
+  createDuckDbTable("u", buildVectors);
+
+  auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
+  auto buildPlanNode =
+      PlanBuilder(planNodeIdGenerator).values(buildVectors).planNode();
+  auto probePlanNode =
+      PlanBuilder(planNodeIdGenerator).values(probeVectors).planNode();
+
+  auto outputType = ROW({"t_0", "u_0"}, {BIGINT(), BIGINT()});
+
+  // Build the HashTable for build side data
+  std::vector<std::unique_ptr<VectorHasher>> hashers;
+  hashers.push_back(std::make_unique<VectorHasher>(BIGINT(), 0));
+
+  std::shared_ptr<facebook::velox::exec::BaseHashTable> table =
+      HashTable<false>::createForJoin(
+          std::move(hashers),
+          {}, /*dependentTypes*/
+          true /*allowDuplicates*/,
+          true /*hasProbedFlag*/,
+          1 /*minTableSizeForParallelJoinBuild*/,
+          pool());
+
+  auto rowContainer = table->rows();
+  uint32_t numColumns = buildVectors[0]->childrenSize();
+  std::vector<DecodedVector> decodedVectors;
+  decodedVectors.reserve(numColumns);
+  for (const auto& rowVector : buildVectors) {
+    if (!rowVector || rowVector->size() == 0)
+      continue;
+
+    decodedVectors.clear();
+    SelectivityVector rows(rowVector->size());
+    for (auto& child : rowVector->children()) {
+      decodedVectors.emplace_back(*child, rows);
+    }
+
+    for (auto i = 0; i < rowVector->size(); ++i) {
+      auto* row = rowContainer->newRow();
+      for (auto j = 0; j < numColumns; ++j) {
+        rowContainer->store(decodedVectors[j], i, row, j);
+      }
+    }
+  }
+
+  table->prepareJoinTable(
+      {}, BaseHashTable::kNoSpillInputStartPartitionBit, 1'000'000);
+
+  auto opaqueSharedHashTable = std::shared_ptr<core::OpaqueHashTable>(
+      table, reinterpret_cast<core::OpaqueHashTable*>(table.get()));
+
+  auto joinNode =
+      core::HashJoinNode::Builder()
+          .id(planNodeIdGenerator->next())
+          .joinType(core::JoinType::kInner)
+          .nullAware(false)
+          .leftKeys(
+              {std::make_shared<core::FieldAccessTypedExpr>(INTEGER(), "t_0")})
+          .rightKeys(
+              {std::make_shared<core::FieldAccessTypedExpr>(INTEGER(), "u_0")})
+          .left(probePlanNode)
+          .right(buildPlanNode)
+          .outputType(outputType)
+          .reusableHashTable(opaqueSharedHashTable)
+          .build();
+
+  std::atomic_bool reusedHashTable{false};
+  SCOPED_TESTVALUE_SET(
+      "facebook::velox::exec::HashBuild::HashBuild",
+      std::function<void(HashBuild*)>(
+          [&](HashBuild* windowBuild) { reusedHashTable.store(true); }));
+
+  auto task =
+      AssertQueryBuilder(joinNode, duckDbQueryRunner_)
+          .maxDrivers(1)
+          .assertResults("SELECT t.t_0, u.u_0 FROM t, u WHERE t.t_0 = u.u_0");
+  ASSERT_TRUE(reusedHashTable.load());
+}
+
 DEBUG_ONLY_TEST_P(
     HashJoinTest,
     hashProbeShouldYieldWhenFilterConsistentlyRejectAll) {