[ut](spill) add UT for partitioned hash join (apache#48432)

be/src/pipeline/exec/hashjoin_build_sink.h

@@ -17,6 +17,7 @@
 
 #pragma once
 
+#include "common/be_mock_util.h"
 #include "exprs/runtime_filter_slots.h"
 #include "join_build_sink_operator.h"
 #include "operator.h"
@@ -25,7 +26,7 @@ namespace doris::pipeline {
 #include "common/compile_check_begin.h"
 class HashJoinBuildSinkOperatorX;
 
-class HashJoinBuildSinkLocalState final
+class HashJoinBuildSinkLocalState MOCK_REMOVE(final)
         : public JoinBuildSinkLocalState<HashJoinSharedState, HashJoinBuildSinkLocalState> {
 public:
     ENABLE_FACTORY_CREATOR(HashJoinBuildSinkLocalState);
@@ -56,7 +57,7 @@ class HashJoinBuildSinkLocalState final
 
     Status disable_runtime_filters(RuntimeState* state);
 
-    [[nodiscard]] size_t get_reserve_mem_size(RuntimeState* state, bool eos);
+    [[nodiscard]] MOCK_FUNCTION size_t get_reserve_mem_size(RuntimeState* state, bool eos);
 
 protected:
     Status _hash_table_init(RuntimeState* state);
@@ -109,7 +110,7 @@ class HashJoinBuildSinkLocalState final
     RuntimeProfile::Counter* _runtime_filter_init_timer = nullptr;
 };
 
-class HashJoinBuildSinkOperatorX final
+class HashJoinBuildSinkOperatorX MOCK_REMOVE(final)
         : public JoinBuildSinkOperatorX<HashJoinBuildSinkLocalState> {
 public:
     HashJoinBuildSinkOperatorX(ObjectPool* pool, int operator_id, int dest_id,
@@ -129,7 +130,7 @@ class HashJoinBuildSinkOperatorX final
 
     [[nodiscard]] size_t get_memory_usage(RuntimeState* state) const;
 
-    std::string get_memory_usage_debug_str(RuntimeState* state) const;
+    MOCK_FUNCTION std::string get_memory_usage_debug_str(RuntimeState* state) const;
 
     bool should_dry_run(RuntimeState* state) override {
         return _is_broadcast_join && !state->get_sink_local_state()

be/src/pipeline/exec/hashjoin_probe_operator.h

@@ -18,6 +18,7 @@
 
 #include <stdint.h>
 
+#include "common/be_mock_util.h"
 #include "common/status.h"
 #include "operator.h"
 #include "pipeline/exec/join_probe_operator.h"
@@ -43,7 +44,7 @@ using HashTableCtxVariants =
                      ProcessHashTableProbe<TJoinOp::NULL_AWARE_LEFT_SEMI_JOIN>>;
 
 class HashJoinProbeOperatorX;
-class HashJoinProbeLocalState final
+class HashJoinProbeLocalState MOCK_REMOVE(final)
         : public JoinProbeLocalState<HashJoinSharedState, HashJoinProbeLocalState> {
 public:
     using Parent = HashJoinProbeOperatorX;
@@ -124,7 +125,8 @@ class HashJoinProbeLocalState final
     RuntimeProfile::Counter* _non_equal_join_conjuncts_timer = nullptr;
 };
 
-class HashJoinProbeOperatorX final : public JoinProbeOperatorX<HashJoinProbeLocalState> {
+class HashJoinProbeOperatorX MOCK_REMOVE(final)
+        : public JoinProbeOperatorX<HashJoinProbeLocalState> {
 public:
     HashJoinProbeOperatorX(ObjectPool* pool, const TPlanNode& tnode, int operator_id,
                            const DescriptorTbl& descs);

be/src/pipeline/exec/operator.h

@@ -284,7 +284,6 @@ class PipelineXSpillLocalState : public PipelineXLocalState<SharedStateArg> {
     Status init(RuntimeState* state, LocalStateInfo& info) override {
         RETURN_IF_ERROR(PipelineXLocalState<SharedStateArg>::init(state, info));
 
-        _spill_total_timer = ADD_TIMER_WITH_LEVEL(Base::profile(), "SpillTotalTime", 1);
         init_spill_read_counters();
 
         return Status::OK();
@@ -317,6 +316,8 @@ class PipelineXSpillLocalState : public PipelineXLocalState<SharedStateArg> {
     }
 
     void init_spill_read_counters() {
+        _spill_total_timer = ADD_TIMER_WITH_LEVEL(Base::profile(), "SpillTotalTime", 1);
+
         // Spill read counters
         _spill_recover_time = ADD_TIMER_WITH_LEVEL(Base::profile(), "SpillRecoverTime", 1);
 
@@ -669,7 +670,11 @@ class PipelineXSpillSinkLocalState : public PipelineXSinkLocalState<SharedStateA
 
     Status init(RuntimeState* state, LocalSinkStateInfo& info) override {
         RETURN_IF_ERROR(Base::init(state, info));
+        init_spill_counters();
+        return Status::OK();
+    }
 
+    void init_spill_counters() {
         _spill_total_timer = ADD_TIMER_WITH_LEVEL(Base::profile(), "SpillTotalTime", 1);
 
         _spill_write_timer = ADD_TIMER_WITH_LEVEL(Base::profile(), "SpillWriteTime", 1);
@@ -696,7 +701,6 @@ class PipelineXSpillSinkLocalState : public PipelineXSinkLocalState<SharedStateA
                 ADD_COUNTER_WITH_LEVEL(Base::profile(), "SpillMaxRowsOfPartition", TUnit::UNIT, 1);
         _spill_min_rows_of_partition =
                 ADD_COUNTER_WITH_LEVEL(Base::profile(), "SpillMinRowsOfPartition", TUnit::UNIT, 1);
-        return Status::OK();
     }
 
     std::vector<Dependency*> dependencies() const override {

be/src/pipeline/exec/partitioned_hash_join_probe_operator.cpp

@@ -58,6 +58,11 @@ Status PartitionedHashJoinProbeLocalState::init(RuntimeState* state, LocalStateI
                                                   "HashJoinProbeSpillDependency", true);
     state->get_task()->add_spill_dependency(_spill_dependency.get());
 
+    init_counters();
+    return Status::OK();
+}
+
+void PartitionedHashJoinProbeLocalState::init_counters() {
     _partition_timer = ADD_TIMER(profile(), "SpillPartitionTime");
     _partition_shuffle_timer = ADD_TIMER(profile(), "SpillPartitionShuffleTime");
     _spill_build_rows = ADD_COUNTER(profile(), "SpillBuildRows", TUnit::UNIT);
@@ -78,7 +83,6 @@ Status PartitionedHashJoinProbeLocalState::init(RuntimeState* state, LocalStateI
             ADD_COUNTER_WITH_LEVEL(profile(), "ProbeBloksBytesInMem", TUnit::BYTES, 1);
     _memory_usage_reserved =
             ADD_COUNTER_WITH_LEVEL(profile(), "MemoryUsageReserved", TUnit::BYTES, 1);
-    return Status::OK();
 }
 
 #define UPDATE_COUNTER_FROM_INNER(name) \
@@ -417,7 +421,7 @@ Status PartitionedHashJoinProbeLocalState::recover_probe_blocks_from_disk(Runtim
             st = spilled_stream->read_next_block_sync(&block, &eos);
             if (!st.ok()) {
                 break;
-            } else {
+            } else if (!block.empty()) {
                 COUNTER_UPDATE(_recovery_probe_rows, block.rows());
                 COUNTER_UPDATE(_recovery_probe_blocks, 1);
                 read_size += block.allocated_bytes();
@@ -786,7 +790,7 @@ size_t PartitionedHashJoinProbeOperatorX::_revocable_mem_size(RuntimeState* stat
 size_t PartitionedHashJoinProbeOperatorX::get_reserve_mem_size(RuntimeState* state) {
     auto& local_state = get_local_state(state);
     const auto need_to_spill = local_state._shared_state->need_to_spill;
-    if (!need_to_spill || !local_state._child_eos) {
+    if (!need_to_spill || local_state._child_eos) {
         return Base::get_reserve_mem_size(state);
     }
 

be/src/pipeline/exec/partitioned_hash_join_probe_operator.h

@@ -20,6 +20,7 @@
 #include <cstdint>
 #include <memory>
 
+#include "common/be_mock_util.h"
 #include "common/status.h"
 #include "operator.h"
 #include "pipeline/dependency.h"
@@ -36,7 +37,7 @@ namespace pipeline {
 
 class PartitionedHashJoinProbeOperatorX;
 
-class PartitionedHashJoinProbeLocalState final
+class PartitionedHashJoinProbeLocalState MOCK_REMOVE(final)
         : public PipelineXSpillLocalState<PartitionedHashJoinSharedState> {
 public:
     using Parent = PartitionedHashJoinProbeOperatorX;
@@ -65,7 +66,9 @@ class PartitionedHashJoinProbeLocalState final
 
     std::string debug_string(int indentation_level = 0) const override;
 
-    void update_profile_from_inner();
+    MOCK_FUNCTION void update_profile_from_inner();
+
+    void init_counters();
 
     friend class PartitionedHashJoinProbeOperatorX;
 

be/src/pipeline/exec/partitioned_hash_join_sink_operator.cpp

@@ -188,7 +188,7 @@ Status PartitionedHashJoinSinkLocalState::_revoke_unpartitioned_block(
                        -(inner_sink_state->_hash_table_memory_usage->value() +
                          inner_sink_state->_build_arena_memory_usage->value()));
     }
-    auto row_desc = p._child->row_desc();
+    const auto& row_desc = p._child->row_desc();
     const auto num_slots = row_desc.num_slots();
     vectorized::Block build_block;
     int64_t block_old_mem = 0;

be/src/pipeline/exec/partitioned_hash_join_sink_operator.h

@@ -21,6 +21,7 @@
 
 #include <atomic>
 
+#include "common/be_mock_util.h"
 #include "common/status.h"
 #include "operator.h"
 #include "pipeline/exec/hashjoin_build_sink.h"
@@ -51,7 +52,7 @@ class PartitionedHashJoinSinkLocalState
     size_t revocable_mem_size(RuntimeState* state) const;
     [[nodiscard]] size_t get_reserve_mem_size(RuntimeState* state, bool eos);
     void update_memory_usage();
-    void update_profile_from_inner();
+    MOCK_FUNCTION void update_profile_from_inner();
 
     Dependency* finishdependency() override;
 
@@ -143,6 +144,9 @@ class PartitionedHashJoinSinkOperatorX
 
 private:
     friend class PartitionedHashJoinSinkLocalState;
+#ifdef BE_TEST
+    friend class PartitionedHashJoinSinkOperatorTest;
+#endif
 
     const TJoinDistributionType::type _join_distribution;
 

be/src/pipeline/exec/spill_utils.h

@@ -138,9 +138,16 @@ class SpillRunnable : public Runnable {
             return;
         }
 
+        Defer set_ready_defer([&] {
+            if (_spill_dependency) {
+                _spill_dependency->set_ready();
+            }
+        });
+
         if (_state->is_cancelled()) {
             return;
         }
+
         auto status = _spill_exec_func();
         if (!status.ok()) {
             ExecEnv::GetInstance()->fragment_mgr()->cancel_query(_state->query_id(), status);
@@ -153,10 +160,6 @@ class SpillRunnable : public Runnable {
                 ExecEnv::GetInstance()->fragment_mgr()->cancel_query(_state->query_id(), status2);
             }
         }
-
-        if (_spill_dependency) {
-            _spill_dependency->set_ready();
-        }
     }
 
 protected:

be/src/runtime/fragment_mgr.h

@@ -31,6 +31,7 @@
 #include <unordered_map>
 #include <vector>
 
+#include "common/be_mock_util.h"
 #include "common/status.h"
 #include "gutil/ref_counted.h"
 #include "http/rest_monitor_iface.h"
@@ -140,7 +141,7 @@ class FragmentMgr : public RestMonitorIface {
                                            std::shared_ptr<pipeline::PipelineFragmentContext>&&);
 
     // Can be used in both version.
-    void cancel_query(const TUniqueId query_id, const Status reason);
+    MOCK_FUNCTION void cancel_query(const TUniqueId query_id, const Status reason);
 
     void cancel_worker();