[Improvement](sort) Free sort blocks if this block is exhausted

Author: Gabriel39

be/src/vec/common/sort/partition_sorter.cpp

@@ -58,7 +58,7 @@ Status PartitionSorter::append_block(Block* input_block) {
     Block sorted_block = VectorizedUtils::create_empty_columnswithtypename(_row_desc);
     DCHECK(input_block->columns() == sorted_block.columns());
     RETURN_IF_ERROR(partial_sort(*input_block, sorted_block));
-    RETURN_IF_ERROR(_state->add_sorted_block(sorted_block));
+    RETURN_IF_ERROR(_state->add_sorted_block(std::move(sorted_block)));
     return Status::OK();
 }
 
@@ -67,46 +67,49 @@ Status PartitionSorter::prepare_for_read() {
     auto& blocks = _state->get_sorted_block();
     auto& priority_queue = _state->get_priority_queue();
     for (auto& block : blocks) {
-        cursors.emplace_back(block, _sort_description);
+        cursors.emplace_back(MergeSortCursorImpl::create_shared(block, _sort_description));
     }
     for (auto& cursor : cursors) {
-        priority_queue.push(MergeSortCursor(&cursor));
+        priority_queue.push(std::move(cursor));
     }
+    cursors.clear();
     return Status::OK();
 }
 
 // have done sorter and get topn records, so could reset those state to init
 void PartitionSorter::reset_sorter_state(RuntimeState* runtime_state) {
-    std::priority_queue<MergeSortBlockCursor> empty_queue;
+    std::priority_queue<std::shared_ptr<MergeSortBlockCursor>> empty_queue;
     std::swap(_block_priority_queue, empty_queue);
     _state = MergeSorterState::create_unique(_row_desc, _offset, _limit, runtime_state, nullptr);
     _previous_row->reset();
 }
 
 Status PartitionSorter::get_next(RuntimeState* state, Block* block, bool* eos) {
-    if (_state->get_sorted_block().empty()) {
+    if (_state->get_sorted_block().empty() && _state->get_priority_queue().empty()) {
+        *eos = true;
+    } else if (_state->get_cursors().size() == 1 && _has_global_limit) {
+        auto& cursor = _state->get_cursors()[0];
+        block->swap(*cursor->block);
+        block->set_num_rows(_partition_inner_limit);
         *eos = true;
     } else {
-        if (_state->get_sorted_block().size() == 1 && _has_global_limit) {
-            auto& sorted_block = _state->get_sorted_block()[0];
-            block->swap(sorted_block);
-            block->set_num_rows(_partition_inner_limit);
-            *eos = true;
-        } else {
-            RETURN_IF_ERROR(partition_sort_read(block, eos, state->batch_size()));
-        }
+        RETURN_IF_ERROR(partition_sort_read(block, eos, state->batch_size()));
     }
     return Status::OK();
 }
 
 Status PartitionSorter::partition_sort_read(Block* output_block, bool* eos, int batch_size) {
-    const auto& sorted_block = _state->get_sorted_block()[0];
-    size_t num_columns = sorted_block.columns();
+    auto& priority_queue = _state->get_priority_queue();
+    if (priority_queue.empty()) {
+        *eos = true;
+        return Status::OK();
+    }
+    const auto& sorted_block = priority_queue.top().impl->block;
+    size_t num_columns = sorted_block->columns();
     MutableBlock m_block =
-            VectorizedUtils::build_mutable_mem_reuse_block(output_block, sorted_block);
+            VectorizedUtils::build_mutable_mem_reuse_block(output_block, *sorted_block);
     MutableColumns& merged_columns = m_block.mutable_columns();
     size_t current_output_rows = 0;
-    auto& priority_queue = _state->get_priority_queue();
 
     bool get_enough_data = false;
     while (!priority_queue.empty()) {
@@ -121,7 +124,7 @@ Status PartitionSorter::partition_sort_read(Block* output_block, bool* eos, int
             //1 row_number no need to check distinct, just output partition_inner_limit row
             if ((current_output_rows + _output_total_rows) < _partition_inner_limit) {
                 for (size_t i = 0; i < num_columns; ++i) {
-                    merged_columns[i]->insert_from(*current->all_columns[i], current->pos);
+                    merged_columns[i]->insert_from(*current->block->get_columns()[i], current->pos);
                 }
             } else {
                 //rows has get enough
@@ -155,7 +158,7 @@ Status PartitionSorter::partition_sort_read(Block* output_block, bool* eos, int
                 }
             }
             for (size_t i = 0; i < num_columns; ++i) {
-                merged_columns[i]->insert_from(*current->all_columns[i], current->pos);
+                merged_columns[i]->insert_from(*current->block->get_columns()[i], current->pos);
             }
             break;
         }
@@ -180,7 +183,7 @@ Status PartitionSorter::partition_sort_read(Block* output_block, bool* eos, int
                 *_previous_row = current;
             }
             for (size_t i = 0; i < num_columns; ++i) {
-                merged_columns[i]->insert_from(*current->all_columns[i], current->pos);
+                merged_columns[i]->insert_from(*current->block->get_columns()[i], current->pos);
             }
             current_output_rows++;
             break;

be/src/vec/common/sort/partition_sorter.h

@@ -50,7 +50,7 @@ struct SortCursorCmp {
     SortCursorCmp(const MergeSortCursor& cursor) : row(cursor->pos), impl(cursor.impl) {}
 
     void reset() {
-        impl = nullptr;
+        impl->reset();
         row = 0;
     }
     bool compare_two_rows(const MergeSortCursor& rhs) const {
@@ -67,7 +67,7 @@ struct SortCursorCmp {
         return true;
     }
     int row = 0;
-    MergeSortCursorImpl* impl = nullptr;
+    std::shared_ptr<MergeSortCursorImpl> impl = nullptr;
 };
 
 class PartitionSorter final : public Sorter {

be/src/vec/common/sort/sorter.cpp

@@ -59,47 +59,50 @@ namespace doris::vectorized {
 void MergeSorterState::reset() {
     auto empty_queue = std::priority_queue<MergeSortCursor>();
     priority_queue_.swap(empty_queue);
-    std::vector<MergeSortCursorImpl> empty_cursors(0);
+    std::vector<std::shared_ptr<MergeSortCursorImpl>> empty_cursors(0);
     cursors_.swap(empty_cursors);
-    std::vector<Block> empty_blocks(0);
+    std::vector<std::shared_ptr<Block>> empty_blocks(0);
     sorted_blocks_.swap(empty_blocks);
     unsorted_block_ = Block::create_unique(unsorted_block_->clone_empty());
     in_mem_sorted_bocks_size_ = 0;
 }
-Status MergeSorterState::add_sorted_block(Block& block) {
+Status MergeSorterState::add_sorted_block(Block&& block) {
     auto rows = block.rows();
     if (0 == rows) {
         return Status::OK();
     }
     in_mem_sorted_bocks_size_ += block.bytes();
-    sorted_blocks_.emplace_back(std::move(block));
+    sorted_blocks_.emplace_back(Block::create_shared(std::move(block)));
     num_rows_ += rows;
     return Status::OK();
 }
 
 Status MergeSorterState::build_merge_tree(const SortDescription& sort_description) {
     for (auto& block : sorted_blocks_) {
-        cursors_.emplace_back(block, sort_description);
+        cursors_.emplace_back(
+                MergeSortCursorImpl::create_shared(std::move(block), sort_description));
     }
 
-    if (sorted_blocks_.size() > 1) {
+    sorted_blocks_.clear();
+    if (cursors_.size() > 1) {
         for (auto& cursor : cursors_) {
-            priority_queue_.emplace(&cursor);
+            priority_queue_.emplace(std::move(cursor));
         }
+        cursors_.clear();
     }
 
     return Status::OK();
 }
 
 Status MergeSorterState::merge_sort_read(doris::vectorized::Block* block, int batch_size,
                                          bool* eos) {
-    if (sorted_blocks_.empty()) {
+    if (cursors_.empty() && priority_queue_.empty()) {
         *eos = true;
-    } else if (sorted_blocks_.size() == 1) {
+    } else if (cursors_.size() == 1) {
         if (offset_ != 0) {
-            sorted_blocks_[0].skip_num_rows(offset_);
+            cursors_[0]->block->skip_num_rows(offset_);
         }
-        block->swap(sorted_blocks_[0]);
+        block->swap(*cursors_[0]->block);
         *eos = true;
     } else {
         RETURN_IF_ERROR(_merge_sort_read_impl(batch_size, block, eos));
@@ -109,9 +112,14 @@ Status MergeSorterState::merge_sort_read(doris::vectorized::Block* block, int ba
 
 Status MergeSorterState::_merge_sort_read_impl(int batch_size, doris::vectorized::Block* block,
                                                bool* eos) {
-    size_t num_columns = sorted_blocks_[0].columns();
+    if (priority_queue_.empty()) {
+        *eos = true;
+        return Status::OK();
+    }
+    size_t num_columns = priority_queue_.top().impl->block->columns();
 
-    MutableBlock m_block = VectorizedUtils::build_mutable_mem_reuse_block(block, sorted_blocks_[0]);
+    MutableBlock m_block = VectorizedUtils::build_mutable_mem_reuse_block(
+            block, *priority_queue_.top().impl->block);
     MutableColumns& merged_columns = m_block.mutable_columns();
 
     /// Take rows from queue in right order and push to 'merged'.
@@ -122,7 +130,7 @@ Status MergeSorterState::_merge_sort_read_impl(int batch_size, doris::vectorized
 
         if (offset_ == 0) {
             for (size_t i = 0; i < num_columns; ++i)
-                merged_columns[i]->insert_from(*current->all_columns[i], current->pos);
+                merged_columns[i]->insert_from(*current->block->get_columns()[i], current->pos);
             ++merged_rows;
         } else {
             offset_--;
@@ -133,7 +141,9 @@ Status MergeSorterState::_merge_sort_read_impl(int batch_size, doris::vectorized
             priority_queue_.push(current);
         }
 
-        if (merged_rows == batch_size) break;
+        if (merged_rows == batch_size) {
+            break;
+        }
     }
     block->set_columns(std::move(merged_columns));
 
@@ -260,22 +270,24 @@ Status FullSorter::_do_sort() {
         // if one block totally greater the heap top of _block_priority_queue
         // we can throw the block data directly.
         if (_state->num_rows() < _offset + _limit) {
-            static_cast<void>(_state->add_sorted_block(desc_block));
-            _block_priority_queue.emplace(_pool->add(
-                    new MergeSortCursorImpl(_state->last_sorted_block(), _sort_description)));
+            static_cast<void>(_state->add_sorted_block(std::move(desc_block)));
+            _block_priority_queue.emplace(
+                    MergeSortBlockCursor::create_shared(MergeSortCursorImpl::create_shared(
+                            _state->last_sorted_block(), _sort_description)));
         } else {
-            auto tmp_cursor_impl =
-                    std::make_unique<MergeSortCursorImpl>(desc_block, _sort_description);
-            MergeSortBlockCursor block_cursor(tmp_cursor_impl.get());
-            if (!block_cursor.totally_greater(_block_priority_queue.top())) {
-                static_cast<void>(_state->add_sorted_block(desc_block));
-                _block_priority_queue.emplace(_pool->add(
-                        new MergeSortCursorImpl(_state->last_sorted_block(), _sort_description)));
+            auto tmp_cursor_impl = MergeSortCursorImpl::create_shared(
+                    Block::create_shared(std::move(desc_block)), _sort_description);
+            MergeSortBlockCursor block_cursor(tmp_cursor_impl);
+            if (!block_cursor.totally_greater(*_block_priority_queue.top())) {
+                static_cast<void>(_state->add_sorted_block(std::move(*tmp_cursor_impl->block)));
+                _block_priority_queue.emplace(
+                        MergeSortBlockCursor::create_shared(MergeSortCursorImpl::create_shared(
+                                _state->last_sorted_block(), _sort_description)));
             }
         }
     } else {
         // dispose normal sort logic
-        static_cast<void>(_state->add_sorted_block(desc_block));
+        static_cast<void>(_state->add_sorted_block(std::move(desc_block)));
     }
     return Status::OK();
 }

be/src/vec/common/sort/sorter.h

@@ -59,7 +59,7 @@ class MergeSorterState {
 
     ~MergeSorterState() = default;
 
-    Status add_sorted_block(Block& block);
+    Status add_sorted_block(Block&& block);
 
     Status build_merge_tree(const SortDescription& sort_description);
 
@@ -72,23 +72,21 @@ class MergeSorterState {
 
     uint64_t num_rows() const { return num_rows_; }
 
-    Block& last_sorted_block() { return sorted_blocks_.back(); }
+    std::shared_ptr<Block> last_sorted_block() { return sorted_blocks_.back(); }
 
-    std::vector<Block>& get_sorted_block() { return sorted_blocks_; }
+    std::vector<std::shared_ptr<Block>>& get_sorted_block() { return sorted_blocks_; }
     std::priority_queue<MergeSortCursor>& get_priority_queue() { return priority_queue_; }
-    std::vector<MergeSortCursorImpl>& get_cursors() { return cursors_; }
+    std::vector<std::shared_ptr<MergeSortCursorImpl>>& get_cursors() { return cursors_; }
     void reset();
 
     std::unique_ptr<Block> unsorted_block_;
 
 private:
-    int _calc_spill_blocks_to_merge() const;
-
     Status _merge_sort_read_impl(int batch_size, doris::vectorized::Block* block, bool* eos);
 
     std::priority_queue<MergeSortCursor> priority_queue_;
-    std::vector<MergeSortCursorImpl> cursors_;
-    std::vector<Block> sorted_blocks_;
+    std::vector<std::shared_ptr<MergeSortCursorImpl>> cursors_;
+    std::vector<std::shared_ptr<Block>> sorted_blocks_;
     size_t in_mem_sorted_bocks_size_ = 0;
     uint64_t num_rows_ = 0;
 
@@ -153,7 +151,7 @@ class Sorter {
     RuntimeProfile::Counter* _partial_sort_timer = nullptr;
     RuntimeProfile::Counter* _merge_block_timer = nullptr;
 
-    std::priority_queue<MergeSortBlockCursor> _block_priority_queue;
+    std::priority_queue<std::shared_ptr<MergeSortBlockCursor>> _block_priority_queue;
     bool _materialize_sort_exprs;
 };
 

be/src/vec/common/sort/topn_sorter.cpp

@@ -72,17 +72,19 @@ Status TopNSorter::_do_sort(Block* block) {
         // if one block totally greater the heap top of _block_priority_queue
         // we can throw the block data directly.
         if (_state->num_rows() < _offset + _limit) {
-            RETURN_IF_ERROR(_state->add_sorted_block(sorted_block));
-            _block_priority_queue.emplace(_pool->add(
-                    new MergeSortCursorImpl(_state->last_sorted_block(), _sort_description)));
+            RETURN_IF_ERROR(_state->add_sorted_block(std::move(sorted_block)));
+            _block_priority_queue.emplace(
+                    MergeSortBlockCursor::create_shared(MergeSortCursorImpl::create_shared(
+                            _state->last_sorted_block(), _sort_description)));
         } else {
-            auto tmp_cursor_impl =
-                    std::make_unique<MergeSortCursorImpl>(sorted_block, _sort_description);
-            MergeSortBlockCursor block_cursor(tmp_cursor_impl.get());
-            if (!block_cursor.totally_greater(_block_priority_queue.top())) {
-                RETURN_IF_ERROR(_state->add_sorted_block(sorted_block));
-                _block_priority_queue.emplace(_pool->add(
-                        new MergeSortCursorImpl(_state->last_sorted_block(), _sort_description)));
+            auto tmp_cursor_impl = MergeSortCursorImpl::create_shared(
+                    Block::create_shared(std::move(sorted_block)), _sort_description);
+            MergeSortBlockCursor block_cursor(tmp_cursor_impl);
+            if (!block_cursor.totally_greater(*_block_priority_queue.top())) {
+                RETURN_IF_ERROR(_state->add_sorted_block(std::move(*block_cursor.impl->block)));
+                _block_priority_queue.emplace(MergeSortBlockCursor::create_shared(
+                        MergeSortCursorImpl::create_shared(MergeSortCursorImpl(
+                                _state->last_sorted_block(), _sort_description))));
             }
         }
     } else {