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

zhangstar333 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 04df4b7decb [fix](be) Preserve rank peer groups across PartitionTopN 
batches (#67894)
04df4b7decb is described below

commit 04df4b7decb9805b1cdc0b0be052089447afa95d
Author: HappenLee <[email protected]>
AuthorDate: Mon Sep 14 20:08:34 2026 +0800

    [fix](be) Preserve rank peer groups across PartitionTopN batches (#67894)
    
    Problem Summary:
    
    Filtering `rank()` or `dense_rank()` through PartitionTopN can silently
    lose rows when the last qualifying peer group spans output batches. With
    `batch_size=4` and five rows sharing the lowest order key, `rank <= 1`
    returns only four rows. The rank reader's deferred EOS check treats a
    satisfied rank counter as the end of the peer group, causing final
    output and intermediate sink pruning to stop too early.
    
    Report EOS only when the next peer group is outside the rank limit or
    the merge queue is exhausted. Preserve the comparison cursor and
    counters across full output batches. Between sorting passes, clear the
    previous-peer comparison cursor along with the counters so the first
    peer group is counted again, avoiding an extra DENSE_RANK group after
    sorter reuse.
    
    After intermediate pruning, set the next fresh-row threshold to the
    larger of the base interval and the retained output row count. This
    reduces repeated processing of growing peer groups and restores the base
    interval when the retained result shrinks. With 16 all-peer input
    batches, the unit test verifies that pruning processes only the prefixes
    of 1, 2, 4, 8, and 16 batches. After recording the next threshold, reset
    the completed sorter immediately to release unread rejected blocks and
    the saved peer cursor before buffering more input. Create the sorter
    only when needed at the next pass or final preparation, avoiding
    duplicate resets. Delaying pruning can temporarily retain more new rows
    that will later be discarded.
    
    Parameterized BE tests cover batch boundaries, multiple merge cursors,
    exhausted input, rank versus dense_rank limits, sorter reuse, adaptive
    pruning, and immediate cursor cleanup after a selective shrink that
    remains above the base threshold. The cleanup test also exercises the
    next pruning pass and final sorter output. SQL regression coverage
    includes global and partitioned windows, a two-row boundary group
    starting mid-batch, ROW_NUMBER, and intermediate pruning with 21,005
    qualifying peers.
    
    ### Release note
    
    Fix missing rows in rank() and dense_rank() filters when PartitionTopN
    peer groups span output batches, reduce repeated processing of large
    peer groups, and release rejected blocks promptly after intermediate
    pruning.
    
    ### Check List (For Author)
    
    - Test:
    - Current head: 19 ASAN unit tests passed in `PartitionSorterTest` and
    `RankAlgorithms/PartitionSorterRankTest` through `run-be-ut.sh -j 48
    --run
    --filter='PartitionSorterTest.*:RankAlgorithms/PartitionSorterRankTest.*'`.
    - Latest cleanup changes: clang-format 16.0.6, build hygiene, and
    clang-tidy passed for all four changed files.
    - Earlier validation at `5c59cc66f89`: BE ASAN build and the
    `test_partition_topn_rank_batch` / `test_partition_topn` SQL regression
    suites passed. These SQL suites were not rerun for the follow-up fixes.
    - Earlier regression oracle: expected output was generated by the runner
    with PartitionTopN disabled and matched by the fixed optimized path. The
    old optimized binary failed `rank_global`, omitting a peer row.
    - Behavior changed: Yes. Preserve complete qualifying rank/dense_rank
    peer groups, correctly reset peer counting between passes, adapt
    intermediate pruning frequency to retained rows, and release completed
    sorter state before buffering the next interval.
    - Does this need documentation: No
---
 be/src/exec/common/partition_sort_utils.cpp        |  18 +-
 be/src/exec/common/partition_sort_utils.h          |   3 +-
 .../exec/operator/partition_sort_sink_operator.cpp |   2 +-
 be/src/exec/sort/partition_sorter.cpp              |  11 +-
 be/src/exec/sort/partition_sorter.h                |   4 +-
 be/test/exec/sort/partition_sorter_test.cpp        | 232 ++++++++++++++++++++-
 .../test_partition_topn_rank_batch.out             | 109 ++++++++++
 .../test_partition_topn_rank_batch.groovy          | 105 ++++++++++
 8 files changed, 468 insertions(+), 16 deletions(-)

diff --git a/be/src/exec/common/partition_sort_utils.cpp 
b/be/src/exec/common/partition_sort_utils.cpp
index ed042b1686d..3318ccf2f41 100644
--- a/be/src/exec/common/partition_sort_utils.cpp
+++ b/be/src/exec/common/partition_sort_utils.cpp
@@ -17,6 +17,8 @@
 
 #include "exec/common/partition_sort_utils.h"
 
+#include <algorithm>
+
 namespace doris {
 
 Status PartitionBlocks::append_block_by_selector(const Block* input_block, 
bool eos) {
@@ -40,19 +42,25 @@ Status PartitionBlocks::append_block_by_selector(const 
Block* input_block, bool
         _init_rows = _init_rows - selector_rows;
         _current_input_rows = _current_input_rows + selector_rows;
         _selector.clear();
-        // maybe better could change by user PARTITION_SORT_ROWS_THRESHOLD
         if (!eos && _partition_sort_info->_partition_inner_limit != -1 &&
-            _current_input_rows >= PARTITION_SORT_ROWS_THRESHOLD &&
+            _current_input_rows >= _partition_sort_rows_threshold &&
             _partition_sort_info->_topn_phase != 
TPartTopNPhase::TWO_PHASE_GLOBAL) {
-            create_or_reset_sorter_state();
+            create_sorter_if_needed();
             RETURN_IF_ERROR(do_partition_topn_sort());
+            // Amortize reprocessing retained peers over at least as many 
fresh rows.
+            _partition_sort_rows_threshold =
+                    std::max(PARTITION_SORT_ROWS_THRESHOLD,
+                             
static_cast<size_t>(_partition_topn_sorter->get_output_rows()));
+            // Retained rows are already in _blocks. Release the completed 
pass's
+            // unread blocks and previous peer cursor before buffering more 
input.
+            
_partition_topn_sorter->reset_sorter_state(_partition_sort_info->_runtime_state);
             _current_input_rows = 0; // reset record
         }
     }
     return Status::OK();
 }
 
-void PartitionBlocks::create_or_reset_sorter_state() {
+void PartitionBlocks::create_sorter_if_needed() {
     if (_partition_topn_sorter == nullptr) {
         _previous_row = std::make_unique<SortCursorCmp>();
         _partition_topn_sorter = PartitionSorter::create_unique(
@@ -65,8 +73,6 @@ void PartitionBlocks::create_or_reset_sorter_state() {
                 _partition_sort_info->_partition_inner_limit,
                 _partition_sort_info->_top_n_algorithm, _previous_row.get());
         
_partition_topn_sorter->init_profile(_partition_sort_info->_runtime_profile);
-    } else {
-        
_partition_topn_sorter->reset_sorter_state(_partition_sort_info->_runtime_state);
     }
 }
 
diff --git a/be/src/exec/common/partition_sort_utils.h 
b/be/src/exec/common/partition_sort_utils.h
index b4cfb0809fd..eb96d17a871 100644
--- a/be/src/exec/common/partition_sort_utils.h
+++ b/be/src/exec/common/partition_sort_utils.h
@@ -90,7 +90,7 @@ public:
 
     Status do_partition_topn_sort();
 
-    void create_or_reset_sorter_state();
+    void create_sorter_if_needed();
 
     void append_whole_block(Block* input_block, const RowDescriptor& row_desc) 
{
         auto empty_block = 
Block::create_unique(VectorizedUtils::create_empty_block(row_desc));
@@ -105,6 +105,7 @@ public:
     IColumn::Selector _selector;
     std::vector<std::unique_ptr<Block>> _blocks;
     size_t _current_input_rows = 0;
+    size_t _partition_sort_rows_threshold = PARTITION_SORT_ROWS_THRESHOLD;
     int64_t _init_rows = 4096;
     bool _is_first_sorter = false;
 
diff --git a/be/src/exec/operator/partition_sort_sink_operator.cpp 
b/be/src/exec/operator/partition_sort_sink_operator.cpp
index 802fdb758ed..e8611ed70a7 100644
--- a/be/src/exec/operator/partition_sort_sink_operator.cpp
+++ b/be/src/exec/operator/partition_sort_sink_operator.cpp
@@ -148,7 +148,7 @@ Status PartitionSortSinkOperatorX::sink_impl(RuntimeState* 
state, Block* input_b
         local_state._partitioned_data.reset(nullptr);
         SCOPED_TIMER(local_state._sorted_data_timer);
         for (auto& _value_place : local_state._value_places) {
-            _value_place->create_or_reset_sorter_state();
+            _value_place->create_sorter_if_needed();
             local_state._shared_state->partition_sorts.emplace_back(
                     std::move(_value_place->_partition_topn_sorter));
         }
diff --git a/be/src/exec/sort/partition_sorter.cpp 
b/be/src/exec/sort/partition_sorter.cpp
index 87b915990d1..23a254d7066 100644
--- a/be/src/exec/sort/partition_sorter.cpp
+++ b/be/src/exec/sort/partition_sorter.cpp
@@ -154,6 +154,7 @@ Status PartitionSorter::_read_row_num(Block* output_block, 
bool* eos, int batch_
 }
 
 Status PartitionSorter::_read_row_rank(Block* output_block, bool* eos, int 
batch_size) {
+    *eos = false;
     auto& queue = _state->get_queue();
     size_t num_columns = _state->unsorted_block()->columns();
 
@@ -163,12 +164,6 @@ Status PartitionSorter::_read_row_rank(Block* 
output_block, bool* eos, int batch
     MutableColumns& merged_columns = m_block.mutable_columns();
     size_t merged_rows = 0;
 
-    Defer defer {[&]() {
-        if (merged_rows == 0 || _get_enough_data()) {
-            *eos = true;
-        }
-    }};
-
     while (queue.is_valid() && merged_rows < batch_size) {
         auto [current, current_rows] = queue.current();
 
@@ -182,6 +177,7 @@ Status PartitionSorter::_read_row_rank(Block* output_block, 
bool* eos, int batch
                 // rank() maybe need check when have get a distinct row
                 // so when the cmp_res is get a distinct row, need check have 
output all rows num
                 if (_get_enough_data()) {
+                    *eos = true;
                     scoped_mutable_block.restore();
                     return Status::OK();
                 }
@@ -201,6 +197,9 @@ Status PartitionSorter::_read_row_rank(Block* output_block, 
bool* eos, int batch
         }
     }
 
+    // A full batch can end inside the last qualifying peer group. Continue 
reading
+    // until the next group exceeds the rank limit or the merge queue is 
exhausted.
+    *eos = !queue.is_valid();
     return Status::OK();
 }
 
diff --git a/be/src/exec/sort/partition_sorter.h 
b/be/src/exec/sort/partition_sorter.h
index 707d992a0d7..0930f611665 100644
--- a/be/src/exec/sort/partition_sorter.h
+++ b/be/src/exec/sort/partition_sorter.h
@@ -48,7 +48,9 @@ public:
     SortCursorCmp(const MergeSortCursor& cursor) : row(cursor->pos), 
impl(cursor.impl) {}
 
     void reset() {
-        impl->reset();
+        // A new sorting pass has no previous peer group. Rewinding the old 
cursor
+        // would make a matching first group appear already seen and go 
uncounted.
+        impl.reset();
         row = 0;
     }
     bool compare_two_rows(const MergeSortCursor& rhs) const {
diff --git a/be/test/exec/sort/partition_sorter_test.cpp 
b/be/test/exec/sort/partition_sorter_test.cpp
index 8bf863e840c..551c51cfba2 100644
--- a/be/test/exec/sort/partition_sorter_test.cpp
+++ b/be/test/exec/sort/partition_sorter_test.cpp
@@ -27,10 +27,12 @@
 #include <memory>
 #include <random>
 #include <utility>
+#include <vector>
 
 #include "common/object_pool.h"
 #include "core/assert_cast.h"
 #include "core/block/block.h"
+#include "exec/common/partition_sort_utils.h"
 #include "exec/sort/heap_sorter.h"
 #include "exec/sort/sorter.h"
 #include "exec/sort/topn_sorter.h"
@@ -187,4 +189,232 @@ TEST_F(PartitionSorterTest, test_partition_sorter_RANK) {
     sorter->reset_sorter_state(&_state);
 }
 
-} // namespace doris
\ No newline at end of file
+struct PartitionSorterRankTest : PartitionSorterTest,
+                                 
testing::WithParamInterface<TopNAlgorithm::type> {
+    std::unique_ptr<PartitionBlocks> create_partition_blocks() {
+        _state._batch_size = PARTITION_SORT_ROWS_THRESHOLD;
+        auto sort_info = std::make_shared<PartitionSortInfo>(
+                &ordering_expr_ctxs, -1, 0, &pool, is_asc_order, nulls_first, 
*row_desc, &_state,
+                &_profile, false, 1, GetParam(), 
TPartTopNPhase::TWO_PHASE_LOCAL);
+        return std::make_unique<PartitionBlocks>(std::move(sort_info), true);
+    }
+
+    void append_rows(PartitionBlocks& partition, const std::vector<int64_t>& 
values) {
+        auto block = ColumnHelper::create_block<DataTypeInt64>(values);
+        for (size_t i = 0; i < values.size(); ++i) {
+            partition.add_row_idx(i);
+        }
+        ASSERT_TRUE(partition.append_block_by_selector(&block, false).ok());
+    }
+
+    void check_retained_rows(const PartitionBlocks& partition, size_t rows, 
int64_t value) {
+        size_t retained_rows = 0;
+        for (const auto& block : partition._blocks) {
+            retained_rows += block->rows();
+            EXPECT_TRUE(ColumnHelper::block_equal(
+                    *block, ColumnHelper::create_block<DataTypeInt64>(
+                                    std::vector<int64_t>(block->rows(), 
value))));
+        }
+        EXPECT_EQ(retained_rows, rows);
+    }
+
+    void check_output(int64_t limit, const std::vector<std::vector<int64_t>>& 
inputs,
+                      const std::vector<int64_t>& expected) {
+        _state._batch_size = 4;
+        SortCursorCmp previous_row;
+        auto rank_sorter = PartitionSorter::create_unique(
+                ordering_expr_ctxs, -1, 0, &pool, is_asc_order, nulls_first, 
*row_desc, &_state,
+                nullptr, false, limit, GetParam(), &previous_row);
+        rank_sorter->init_profile(&_profile);
+        for (const auto& values : inputs) {
+            auto block = ColumnHelper::create_block<DataTypeInt64>(values);
+            ASSERT_TRUE(rank_sorter->append_block(&block).ok());
+        }
+        ASSERT_TRUE(rank_sorter->prepare_for_read(false).ok());
+
+        bool eos = false;
+        size_t output_rows = 0;
+        Block block;
+        // Allow one final empty batch when the next peer group starts at a 
batch boundary.
+        for (size_t batch = 0; !eos && batch <= expected.size() / 
_state.batch_size() + 1;
+             ++batch) {
+            block.clear_column_data();
+            ASSERT_TRUE(rank_sorter->get_next(&_state, &block, &eos).ok());
+            const auto rows = std::min<size_t>(_state.batch_size(), 
expected.size() - output_rows);
+            ASSERT_EQ(block.rows(), rows);
+            if (rows > 0) {
+                const std::vector<int64_t> expected_batch(expected.begin() + 
output_rows,
+                                                          expected.begin() + 
output_rows + rows);
+                EXPECT_TRUE(ColumnHelper::block_equal(
+                        block, 
ColumnHelper::create_block<DataTypeInt64>(expected_batch)));
+            }
+            output_rows += rows;
+            if (output_rows < expected.size()) {
+                ASSERT_FALSE(eos);
+            }
+        }
+        EXPECT_TRUE(eos);
+        EXPECT_EQ(output_rows, expected.size());
+    }
+};
+
+TEST_P(PartitionSorterRankTest, BoundaryPeersAcrossBatches) {
+    for (int peer_rows : {3, 4, 5, 9}) {
+        for (bool has_next_group : {false, true}) {
+            SCOPED_TRACE(testing::Message()
+                         << "peer_rows=" << peer_rows << ", has_next_group=" 
<< has_next_group);
+            // Split the peer group between input blocks to also exercise 
merge cursor changes.
+            std::vector<std::vector<int64_t>> inputs {{0}, 
std::vector<int64_t>(peer_rows - 1, 0)};
+            if (has_next_group) {
+                inputs.front().push_back(1);
+            }
+            check_output(1, inputs, std::vector<int64_t>(peer_rows, 0));
+        }
+    }
+}
+
+TEST_P(PartitionSorterRankTest, RankLimitBeyondFirstGroup) {
+    const std::vector<std::vector<int64_t>> inputs {{0, 1, 1, 2}, {0, 1, 1, 1, 
2}};
+    check_output(2, inputs,
+                 GetParam() == TopNAlgorithm::RANK ? std::vector<int64_t> {0, 
0}
+                                                   : std::vector<int64_t> {0, 
0, 1, 1, 1, 1, 1});
+    check_output(GetParam() == TopNAlgorithm::RANK ? 3 : 2, inputs, {0, 0, 1, 
1, 1, 1, 1});
+}
+
+TEST_P(PartitionSorterRankTest, ShortBoundaryGroupAcrossBatches) {
+    // The boundary group can span batches even when it is smaller than a 
batch.
+    check_output(GetParam() == TopNAlgorithm::RANK ? 4 : 2, {{0, 0, 1}, {0, 1, 
2}},
+                 {0, 0, 0, 1, 1});
+}
+
+TEST_P(PartitionSorterRankTest, ExhaustInputBelowLimit) {
+    check_output(10, {{0, 1, 2}, {0, 1}}, {0, 0, 1, 1, 2});
+    check_output(10, {}, {});
+}
+
+TEST_P(PartitionSorterRankTest, ResetStartsNewPeerGroup) {
+    _state._batch_size = 4;
+    SortCursorCmp previous_row;
+    auto rank_sorter = PartitionSorter::create_unique(ordering_expr_ctxs, -1, 
0, &pool,
+                                                      is_asc_order, 
nulls_first, *row_desc, &_state,
+                                                      nullptr, false, 1, 
GetParam(), &previous_row);
+    rank_sorter->init_profile(&_profile);
+    for (int pass = 0; pass < 2; ++pass) {
+        SCOPED_TRACE(pass);
+        auto peers = 
ColumnHelper::create_block<DataTypeInt64>(std::vector<int64_t>(5, 0));
+        auto next_group = 
ColumnHelper::create_block<DataTypeInt64>(std::vector<int64_t>(9, 1));
+        ASSERT_TRUE(rank_sorter->append_block(&peers).ok());
+        ASSERT_TRUE(rank_sorter->append_block(&next_group).ok());
+        ASSERT_TRUE(rank_sorter->prepare_for_read(false).ok());
+
+        bool eos = false;
+        Block output;
+        ASSERT_TRUE(rank_sorter->get_next(&_state, &output, &eos).ok());
+        EXPECT_TRUE(ColumnHelper::block_equal(
+                output, ColumnHelper::create_block<DataTypeInt64>({0, 0, 0, 
0})));
+        ASSERT_FALSE(eos);
+        output.clear_column_data();
+        ASSERT_TRUE(rank_sorter->get_next(&_state, &output, &eos).ok());
+        EXPECT_TRUE(
+                ColumnHelper::block_equal(output, 
ColumnHelper::create_block<DataTypeInt64>({0})));
+        EXPECT_TRUE(eos);
+        rank_sorter->reset_sorter_state(&_state);
+    }
+}
+
+TEST_P(PartitionSorterRankTest, IntermediatePruningAmortizesRetainedPeers) {
+    auto partition = create_partition_blocks();
+    const size_t input_batches = 16;
+    const size_t batch_rows = PARTITION_SORT_ROWS_THRESHOLD;
+    size_t pruning_passes = 0;
+    size_t processed_rows = 0;
+    for (size_t batch = 1; batch <= input_batches; ++batch) {
+        ASSERT_NO_FATAL_FAILURE(append_rows(*partition, 
std::vector<int64_t>(batch_rows, 0)));
+        if (partition->_current_input_rows == 0) {
+            ++pruning_passes;
+            processed_rows += batch * batch_rows;
+        }
+    }
+    // With no rows pruned, only the prefixes of 1, 2, 4, 8 and 16 batches are 
processed.
+    EXPECT_EQ(pruning_passes, 5);
+    EXPECT_LT(processed_rows, 2 * input_batches * batch_rows);
+    check_retained_rows(*partition, input_batches * batch_rows, 0);
+}
+
+TEST_P(PartitionSorterRankTest, IntermediatePruningResumesAfterPeersShrink) {
+    auto partition = create_partition_blocks();
+    const size_t batch_rows = PARTITION_SORT_ROWS_THRESHOLD;
+    for (size_t batch = 0; batch < 4; ++batch) {
+        ASSERT_NO_FATAL_FAILURE(append_rows(*partition, 
std::vector<int64_t>(batch_rows, 0)));
+    }
+    check_retained_rows(*partition, 4 * batch_rows, 0);
+
+    // A better key eliminates the retained peer group on the next pruning 
pass.
+    auto values = std::vector<int64_t>(batch_rows, 0);
+    values.front() = -1;
+    ASSERT_NO_FATAL_FAILURE(append_rows(*partition, values));
+    for (size_t batch = 1; batch < 4; ++batch) {
+        ASSERT_NO_FATAL_FAILURE(append_rows(*partition, 
std::vector<int64_t>(batch_rows, 0)));
+    }
+    check_retained_rows(*partition, 1, -1);
+
+    // Once few rows are retained, one base interval must trigger pruning 
again.
+    ASSERT_NO_FATAL_FAILURE(append_rows(*partition, 
std::vector<int64_t>(batch_rows - 1, -2)));
+    EXPECT_EQ(partition->_current_input_rows, batch_rows - 1);
+    ASSERT_NO_FATAL_FAILURE(append_rows(*partition, {-2}));
+    EXPECT_EQ(partition->_current_input_rows, 0);
+    check_retained_rows(*partition, batch_rows, -2);
+}
+
+TEST_P(PartitionSorterRankTest, IntermediatePruningReleasesRejectedBlocks) {
+    auto partition = create_partition_blocks();
+    const size_t batch_rows = PARTITION_SORT_ROWS_THRESHOLD;
+    for (size_t batch = 0; batch < 4; ++batch) {
+        ASSERT_NO_FATAL_FAILURE(append_rows(*partition, 
std::vector<int64_t>(batch_rows, 0)));
+    }
+
+    // Shrink the retained group from four to two batches, still above the 
base interval.
+    for (size_t batch = 0; batch < 4; ++batch) {
+        ASSERT_NO_FATAL_FAILURE(
+                append_rows(*partition, std::vector<int64_t>(batch_rows, batch 
< 2 ? -1 : 0)));
+    }
+    check_retained_rows(*partition, 2 * batch_rows, -1);
+    EXPECT_EQ(partition->_partition_sort_rows_threshold, 2 * batch_rows);
+    auto& rank_sorter = partition->_partition_topn_sorter;
+    // Early EOS leaves six rejected batches unread. Release the queue and 
saved peer cursor.
+    EXPECT_FALSE(rank_sorter->_state->get_queue().is_valid());
+    EXPECT_EQ(partition->_previous_row->impl, nullptr);
+
+    ASSERT_NO_FATAL_FAILURE(append_rows(*partition, std::vector<int64_t>(2 * 
batch_rows - 1, -2)));
+    EXPECT_EQ(partition->_current_input_rows, 2 * batch_rows - 1);
+    ASSERT_NO_FATAL_FAILURE(append_rows(*partition, {-2}));
+    EXPECT_EQ(partition->_current_input_rows, 0);
+    check_retained_rows(*partition, 2 * batch_rows, -2);
+    EXPECT_FALSE(rank_sorter->_state->get_queue().is_valid());
+    EXPECT_EQ(partition->_previous_row->impl, nullptr);
+
+    // Exercise the final sink preparation with the sorter already reset by 
pruning.
+    partition->create_sorter_if_needed();
+    for (const auto& block : partition->_blocks) {
+        ASSERT_TRUE(rank_sorter->append_block(block.get()).ok());
+    }
+    partition->_blocks.clear();
+    ASSERT_TRUE(rank_sorter->prepare_for_read(false).ok());
+    bool eos = false;
+    size_t output_rows = 0;
+    for (size_t batch = 0; !eos && batch < 3; ++batch) {
+        Block output;
+        ASSERT_TRUE(rank_sorter->get_next(&_state, &output, &eos).ok());
+        output_rows += output.rows();
+        EXPECT_TRUE(ColumnHelper::block_equal(output,
+                                              
ColumnHelper::create_block<DataTypeInt64>(
+                                                      
std::vector<int64_t>(output.rows(), -2))));
+    }
+    EXPECT_TRUE(eos);
+    EXPECT_EQ(output_rows, 2 * batch_rows);
+}
+
+INSTANTIATE_TEST_SUITE_P(RankAlgorithms, PartitionSorterRankTest,
+                         testing::Values(TopNAlgorithm::RANK, 
TopNAlgorithm::DENSE_RANK));
+
+} // namespace doris
diff --git 
a/regression-test/data/query_p0/partition_topn/test_partition_topn_rank_batch.out
 
b/regression-test/data/query_p0/partition_topn/test_partition_topn_rank_batch.out
new file mode 100644
index 00000000000..f0a380ed29b
--- /dev/null
+++ 
b/regression-test/data/query_p0/partition_topn/test_partition_topn_rank_batch.out
@@ -0,0 +1,109 @@
+-- This file is automatically generated. You should know what you did if you 
want to edit this
+-- !rank_global --
+1      0       1
+2      0       1
+3      0       1
+4      0       1
+5      0       1
+
+-- !rank_partitioned --
+1      0       0       1
+2      0       0       1
+3      0       0       1
+4      0       0       1
+5      0       0       1
+7      1       0       1
+8      1       0       1
+9      1       0       1
+10     1       0       1
+11     1       0       1
+12     1       0       1
+13     1       0       1
+14     1       0       1
+15     1       0       1
+
+-- !rank_second_rank --
+1      0       0       1
+2      0       0       1
+3      0       0       1
+4      0       0       1
+5      0       0       1
+7      1       0       1
+8      1       0       1
+9      1       0       1
+10     1       0       1
+11     1       0       1
+12     1       0       1
+13     1       0       1
+14     1       0       1
+15     1       0       1
+
+-- !dense_rank_global --
+1      0       1
+2      0       1
+3      0       1
+4      0       1
+5      0       1
+
+-- !dense_rank_partitioned --
+1      0       0       1
+2      0       0       1
+3      0       0       1
+4      0       0       1
+5      0       0       1
+7      1       0       1
+8      1       0       1
+9      1       0       1
+10     1       0       1
+11     1       0       1
+12     1       0       1
+13     1       0       1
+14     1       0       1
+15     1       0       1
+
+-- !dense_rank_second_rank --
+1      0       0       1
+2      0       0       1
+3      0       0       1
+4      0       0       1
+5      0       0       1
+6      0       1       2
+7      1       0       1
+8      1       0       1
+9      1       0       1
+10     1       0       1
+11     1       0       1
+12     1       0       1
+13     1       0       1
+14     1       0       1
+15     1       0       1
+16     1       1       2
+
+-- !row_number --
+0      0       1
+1      0       1
+
+-- !rank_intermediate_pruning --
+0      5
+1      9
+2      21005
+
+-- !dense_rank_intermediate_pruning --
+0      5
+1      9
+2      21005
+
+-- !rank_short_boundary_group --
+30000  0       1
+30001  0       1
+30002  0       1
+30003  1       4
+30004  1       4
+
+-- !dense_rank_short_boundary_group --
+30000  0       1
+30001  0       1
+30002  0       1
+30003  1       2
+30004  1       2
+
diff --git 
a/regression-test/suites/query_p0/partition_topn/test_partition_topn_rank_batch.groovy
 
b/regression-test/suites/query_p0/partition_topn/test_partition_topn_rank_batch.groovy
new file mode 100644
index 00000000000..da8c6c7a9b5
--- /dev/null
+++ 
b/regression-test/suites/query_p0/partition_topn/test_partition_topn_rank_batch.groovy
@@ -0,0 +1,105 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+suite("test_partition_topn_rank_batch") {
+    sql "set batch_size = 4"
+    sql "set parallel_pipeline_task_num = 1"
+    sql "set enable_partition_topn = true"
+
+    sql "drop table if exists test_partition_topn_rank_batch"
+    sql """
+        create table test_partition_topn_rank_batch (id int, p int, k int)
+        duplicate key(id)
+        distributed by hash(id) buckets 1
+        properties("replication_num" = "1")
+    """
+    sql """
+        insert into test_partition_topn_rank_batch values
+        (1,0,0),(2,0,0),(3,0,0),(4,0,0),(5,0,0),(6,0,1),
+        (7,1,0),(8,1,0),(9,1,0),(10,1,0),(11,1,0),
+        (12,1,0),(13,1,0),(14,1,0),(15,1,0),(16,1,1)
+    """
+
+    for (func in ["rank", "dense_rank"]) {
+        def globalQuery = """
+            select id, k, r from (
+                select id, k, ${func}() over(order by k) r
+                from test_partition_topn_rank_batch where p = 0
+            ) t where r <= 1 order by id
+        """
+        explain {
+            sql globalQuery
+            contains "VPartitionTopN"
+        }
+        "qt_${func}_global"(globalQuery)
+        "qt_${func}_partitioned"("""
+            select id, p, k, r from (
+                select id, p, k, ${func}() over(partition by p order by k) r
+                from test_partition_topn_rank_batch
+            ) t where r <= 1 order by id
+        """)
+        "qt_${func}_second_rank"("""
+            select id, p, k, r from (
+                select id, p, k, ${func}() over(partition by p order by k) r
+                from test_partition_topn_rank_batch
+            ) t where r <= 2 order by id
+        """)
+    }
+
+    qt_row_number """
+        select p, k, r from (
+            select p, k, row_number() over(partition by p order by k) r
+            from test_partition_topn_rank_batch
+        ) t where r <= 1 order by p
+    """
+
+    // Keep more than 20,000 peers before the final input batch to exercise 
intermediate pruning.
+    sql "set batch_size = 1000"
+    sql """
+        insert into test_partition_topn_rank_batch
+        select number + 100, 2, if(number < 21005, 0, 1) from numbers("number" 
= "21006")
+    """
+    for (func in ["rank", "dense_rank"]) {
+        "qt_${func}_intermediate_pruning"("""
+            select p, count(*) from (
+                select p, ${func}() over(partition by p order by k) r
+                from test_partition_topn_rank_batch
+            ) t where r <= 1 group by p order by p
+        """)
+    }
+
+    // A two-row boundary group also spans batches when three rows precede it.
+    sql "set batch_size = 4"
+    sql """
+        insert into test_partition_topn_rank_batch values
+        (30000,3,0),(30001,3,0),(30002,3,0),(30003,3,1),(30004,3,1),(30005,3,2)
+    """
+    for (func in ["rank", "dense_rank"]) {
+        def rankLimit = func == "rank" ? 4 : 2
+        def shortGroupQuery = """
+            select id, k, r from (
+                select id, k, ${func}() over(order by k) r
+                from test_partition_topn_rank_batch where p = 3
+            ) t where r <= ${rankLimit} order by id
+        """
+        explain {
+            sql shortGroupQuery
+            contains "VPartitionTopN"
+        }
+        "qt_${func}_short_boundary_group"(shortGroupQuery)
+    }
+}


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to