(doris) branch master updated: [fix](spill) Fix revocable_mem_size to return 0 when partition reaches max repartition depth (#61973)

[mr...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

mrhhsg pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git


The following commit(s) were added to refs/heads/master by this push:
     new e4ed81fdb1b [fix](spill) Fix revocable_mem_size to return 0 when partition reaches max repartition depth (#61973)
e4ed81fdb1b is described below

commit e4ed81fdb1b78dbf4b283c8fd3096dcf379655b8
Author: Jerry Hu <hu...@selectdb.com>
AuthorDate: Fri Apr 3 10:42:33 2026 +0800

    [fix](spill) Fix revocable_mem_size to return 0 when partition reaches max repartition depth (#61973)
    
    ### What problem does this PR solve?
    
    When a partition's level reaches or exceeds the max repartition depth,
    it cannot be repartitioned further. In this case, `revocable_mem_size()`
    should return 0 because the memory is not actually revocable —
    attempting to spill it would lead to futile repartition attempts.
    
    This fix adds the depth check to both `PartitionedAggSourceOperatorX`
    and `PartitionedHashJoinProbeOperatorX`, along with corresponding unit
    tests that verify:
    - revocable_mem_size returns 0 at (max depth - 1)
    - revocable_mem_size returns >0 at (max depth - 2)
    
    
    
    Issue Number: close #xxx
    
    Related PR: #xxx
    
    Problem Summary:
    
    ### Release note
    
    None
    
    ### Check List (For Author)
    
    - Test <!-- At least one of them must be included. -->
        - [ ] Regression test
        - [ ] Unit Test
        - [ ] Manual test (add detailed scripts or steps below)
        - [ ] No need to test or manual test. Explain why:
    - [ ] This is a refactor/code format and no logic has been changed.
            - [ ] Previous test can cover this change.
            - [ ] No code files have been changed.
            - [ ] Other reason <!-- Add your reason?  -->
    
    - Behavior changed:
        - [ ] No.
        - [ ] Yes. <!-- Explain the behavior change -->
    
    - Does this need documentation?
        - [ ] No.
    - [ ] Yes. <!-- Add document PR link here. eg:
    https://github.com/apache/doris-website/pull/1214 -->
    
    ### Check List (For Reviewer who merge this PR)
    
    - [ ] Confirm the release note
    - [ ] Confirm test cases
    - [ ] Confirm document
    - [ ] Add branch pick label <!-- Add branch pick label that this PR
    should merge into -->
---
 .../partitioned_aggregation_source_operator.cpp    |  5 +++
 .../partitioned_hash_join_probe_operator.cpp       |  5 +++
 ...artitioned_aggregation_source_operator_test.cpp | 38 ++++++++++++++++++
 .../partitioned_hash_join_probe_operator_test.cpp  | 45 ++++++++++++++++++++++
 4 files changed, 93 insertions(+)

diff --git a/be/src/exec/operator/partitioned_aggregation_source_operator.cpp b/be/src/exec/operator/partitioned_aggregation_source_operator.cpp
index bfefcfb9af2..ea57cc0091d 100644
--- a/be/src/exec/operator/partitioned_aggregation_source_operator.cpp
+++ b/be/src/exec/operator/partitioned_aggregation_source_operator.cpp
@@ -191,6 +191,11 @@ size_t PartitionedAggSourceOperatorX::revocable_mem_size(RuntimeState* state) co
     if (!local_state._shared_state->_is_spilled || !local_state._current_partition.spill_file) {
         return 0;
     }
+    // If the current partition has reached the max repartition depth, it cannot be
+    // repartitioned further, so its data is not revocable.
+    if ((local_state._current_partition.level + 1) >= _repartition_max_depth) {
+        return 0;
+    }
 
     size_t bytes = 0;
     for (const auto& block : local_state._blocks) {
diff --git a/be/src/exec/operator/partitioned_hash_join_probe_operator.cpp b/be/src/exec/operator/partitioned_hash_join_probe_operator.cpp
index d5267969cd4..1bbc1972bec 100644
--- a/be/src/exec/operator/partitioned_hash_join_probe_operator.cpp
+++ b/be/src/exec/operator/partitioned_hash_join_probe_operator.cpp
@@ -875,6 +875,11 @@ size_t PartitionedHashJoinProbeOperatorX::revocable_mem_size(RuntimeState* state
         // Or if current partition has finished build hash table.
         return 0;
     }
+    // If the current partition has reached the max repartition depth, it cannot be
+    // repartitioned further, so its data is not revocable.
+    if ((local_state._current_partition.level + 1) >= _repartition_max_depth) {
+        return 0;
+    }
 
     // Include build-side memory that has been recovered but not yet consumed by the hash table.
     // This data is revocable because we can repartition instead of building the hash table.
diff --git a/be/test/exec/operator/partitioned_aggregation_source_operator_test.cpp b/be/test/exec/operator/partitioned_aggregation_source_operator_test.cpp
index f331b529ce5..2b13ad97f4c 100644
--- a/be/test/exec/operator/partitioned_aggregation_source_operator_test.cpp
+++ b/be/test/exec/operator/partitioned_aggregation_source_operator_test.cpp
@@ -737,6 +737,44 @@ TEST_F(PartitionedAggregationSourceOperatorTest, RevocableMemSizeWithAggContaine
     EXPECT_EQ(source_operator->revocable_mem_size(_helper.runtime_state.get()), container_bytes);
 }
 
+// Condition: partition (level + 1) >= _repartition_max_depth → 0 (not revocable).
+TEST_F(PartitionedAggregationSourceOperatorTest, RevocableMemSizeAtMaxDepthReturnsZero) {
+    auto [source_operator, sink_operator] = _helper.create_operators();
+    const auto tnode = _helper.create_test_plan_node();
+    ASSERT_TRUE(source_operator->init(tnode, _helper.runtime_state.get()).ok());
+    ASSERT_TRUE(source_operator->prepare(_helper.runtime_state.get()).ok());
+
+    std::shared_ptr<MockPartitionedAggSharedState> shared_state;
+    auto* local_state = _helper.create_source_local_state(_helper.runtime_state.get(),
+                                                          source_operator.get(), shared_state);
+
+    // Add a large block so that without the depth check, revocable_mem_size would be > 0.
+    std::vector<int32_t> large_data(300000);
+    std::iota(large_data.begin(), large_data.end(), 0);
+    auto large_block = ColumnHelper::create_block<DataTypeInt32>(large_data);
+    ASSERT_GT(large_block.allocated_bytes(), 1UL << 20);
+    local_state->_blocks.push_back(std::move(large_block));
+
+    SpillFileSPtr spill_file;
+    ASSERT_TRUE(ExecEnv::GetInstance()
+                        ->spill_file_mgr()
+                        ->create_spill_file("ut/revocable_max_depth", spill_file)
+                        .ok());
+    local_state->_current_partition.spill_file = spill_file;
+
+    ASSERT_GE(source_operator->_repartition_max_depth, 2);
+
+    // Set partition level to max depth → not revocable.
+    local_state->_current_partition.level =
+            static_cast<int>(source_operator->_repartition_max_depth) - 1;
+    EXPECT_EQ(source_operator->revocable_mem_size(_helper.runtime_state.get()), 0UL);
+
+    // Also verify level == max_depth - 2 IS still revocable.
+    local_state->_current_partition.level =
+            static_cast<int>(source_operator->_repartition_max_depth) - 2;
+    EXPECT_GT(source_operator->revocable_mem_size(_helper.runtime_state.get()), 0UL);
+}
+
 // --- Tests for PartitionedAggSourceOperatorX::revoke_memory ---
 // revoke_memory:
 //   if (!_is_spilled) return OK  (no-op)
diff --git a/be/test/exec/operator/partitioned_hash_join_probe_operator_test.cpp b/be/test/exec/operator/partitioned_hash_join_probe_operator_test.cpp
index 7a1fb8c839a..e50ccb1ec35 100644
--- a/be/test/exec/operator/partitioned_hash_join_probe_operator_test.cpp
+++ b/be/test/exec/operator/partitioned_hash_join_probe_operator_test.cpp
@@ -1027,6 +1027,51 @@ TEST_F(PartitionedHashJoinProbeOperatorTest, revocable_mem_size) {
     ASSERT_EQ(probe_operator->revocable_mem_size(_helper.runtime_state.get()), 0);
 }
 
+// Partition (level + 1) >= _repartition_max_depth → revocable_mem_size returns 0.
+TEST_F(PartitionedHashJoinProbeOperatorTest, revocable_mem_size_at_max_depth) {
+    auto [probe_operator, sink_operator] = _helper.create_operators();
+
+    std::shared_ptr<MockPartitionedHashJoinSharedState> shared_state;
+    auto local_state = _helper.create_probe_local_state(_helper.runtime_state.get(),
+                                                        probe_operator.get(), shared_state);
+
+    local_state->_shared_state->_is_spilled = true;
+    local_state->_child_eos = true;
+
+    // Set up a valid current partition with build not finished (the path that
+    // checks level vs max depth).
+    SpillFileSPtr build_file;
+    ASSERT_TRUE(ExecEnv::GetInstance()
+                        ->spill_file_mgr()
+                        ->create_spill_file("ut/revocable_join_max_depth_build", build_file)
+                        .ok());
+    SpillFileSPtr probe_file;
+    ASSERT_TRUE(ExecEnv::GetInstance()
+                        ->spill_file_mgr()
+                        ->create_spill_file("ut/revocable_join_max_depth_probe", probe_file)
+                        .ok());
+    local_state->_current_partition = JoinSpillPartitionInfo(
+            build_file, probe_file, static_cast<int>(probe_operator->_repartition_max_depth - 1));
+    local_state->_current_partition.build_finished = false;
+
+    // Add a recovered build block large enough to exceed the threshold.
+    std::vector<int32_t> large_data(256 * 1024); // 1MB of int32
+    std::iota(large_data.begin(), large_data.end(), 0);
+    Block large_block = ColumnHelper::create_block<DataTypeInt32>(large_data);
+    ASSERT_GE(large_block.allocated_bytes(), SpillFile::MIN_SPILL_WRITE_BATCH_MEM);
+    local_state->_recovered_build_block = MutableBlock::create_unique(std::move(large_block));
+
+    ASSERT_GE(probe_operator->_repartition_max_depth, 2);
+
+    // At max depth → not revocable.
+    ASSERT_EQ(probe_operator->revocable_mem_size(_helper.runtime_state.get()), 0);
+
+    // One level below max depth → revocable.
+    local_state->_current_partition.level =
+            static_cast<int>(probe_operator->_repartition_max_depth) - 2;
+    ASSERT_GT(probe_operator->revocable_mem_size(_helper.runtime_state.get()), 0);
+}
+
 TEST_F(PartitionedHashJoinProbeOperatorTest, get_reserve_mem_size) {
     // Setup test environment
     auto [probe_operator, sink_operator] = _helper.create_operators();


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@doris.apache.org
For additional commands, e-mail: commits-help@doris.apache.org