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

Gabriel39 · web-flow · commit 44aa31458193 · 2024-08-20T13:13:10.000+08:00
diff --git a/be/src/pipeline/dependency.h b/be/src/pipeline/dependency.h
@@ -886,8 +886,6 @@ struct LocalMergeExchangeSharedState : public LocalExchangeSharedState {
 
     void create_dependencies(int local_exchange_id) override {
         sink_deps.resize(source_deps.size());
-        std::vector<DependencySPtr> new_deps(sink_deps.size(), nullptr);
-        source_deps.swap(new_deps);
         for (size_t i = 0; i < source_deps.size(); i++) {
             source_deps[i] =
                     std::make_shared<Dependency>(local_exchange_id, local_exchange_id,
diff --git a/be/src/vec/common/sort/partition_sorter.cpp b/be/src/vec/common/sort/partition_sorter.cpp
@@ -58,20 +58,17 @@ 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));
+    _state->add_sorted_block(Block::create_shared(std::move(sorted_block)));
     return Status::OK();
 }
 
 Status PartitionSorter::prepare_for_read() {
-    auto& cursors = _state->get_cursors();
     auto& blocks = _state->get_sorted_block();
     auto& priority_queue = _state->get_priority_queue();
     for (auto& block : blocks) {
-        cursors.emplace_back(block, _sort_description);
-    }
-    for (auto& cursor : cursors) {
-        priority_queue.push(MergeSortCursor(&cursor));
+        priority_queue.push(MergeSortCursorImpl::create_shared(block, _sort_description));
     }
+    blocks.clear();
     return Status::OK();
 }
 
@@ -84,29 +81,30 @@ void PartitionSorter::reset_sorter_state(RuntimeState* runtime_state) {
 }
 
 Status PartitionSorter::get_next(RuntimeState* state, Block* block, bool* eos) {
-    if (_state->get_sorted_block().empty()) {
+    if (_state->get_priority_queue().empty()) {
+        *eos = true;
+    } else if (_state->get_priority_queue().size() == 1 && _has_global_limit) {
+        block->swap(*_state->get_priority_queue().top().impl->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 +119,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 +153,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 +178,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;
diff --git a/be/src/vec/common/sort/partition_sorter.h b/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 {
diff --git a/be/src/vec/common/sort/sorter.cpp b/be/src/vec/common/sort/sorter.cpp
@@ -59,48 +59,44 @@ namespace doris::vectorized {
 void MergeSorterState::reset() {
     auto empty_queue = std::priority_queue<MergeSortCursor>();
     priority_queue_.swap(empty_queue);
-    std::vector<MergeSortCursorImpl> empty_cursors(0);
-    cursors_.swap(empty_cursors);
-    std::vector<Block> empty_blocks(0);
+    std::vector<std::shared_ptr<MergeSortCursorImpl>> empty_cursors(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) {
-    auto rows = block.rows();
+void MergeSorterState::add_sorted_block(std::shared_ptr<Block> block) {
+    auto rows = block->rows();
     if (0 == rows) {
-        return Status::OK();
+        return;
     }
-    in_mem_sorted_bocks_size_ += block.bytes();
-    sorted_blocks_.emplace_back(std::move(block));
+    in_mem_sorted_bocks_size_ += block->bytes();
+    sorted_blocks_.emplace_back(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);
-    }
-
-    if (sorted_blocks_.size() > 1) {
-        for (auto& cursor : cursors_) {
-            priority_queue_.emplace(&cursor);
-        }
+        priority_queue_.emplace(
+                MergeSortCursorImpl::create_shared(std::move(block), sort_description));
     }
 
+    sorted_blocks_.clear();
     return Status::OK();
 }
 
 Status MergeSorterState::merge_sort_read(doris::vectorized::Block* block, int batch_size,
                                          bool* eos) {
-    if (sorted_blocks_.empty()) {
+    DCHECK(sorted_blocks_.empty());
+    DCHECK(unsorted_block_->empty());
+    if (priority_queue_.empty()) {
         *eos = true;
-    } else if (sorted_blocks_.size() == 1) {
+    } else if (priority_queue_.size() == 1) {
         if (offset_ != 0) {
-            sorted_blocks_[0].skip_num_rows(offset_);
+            priority_queue_.top().impl->block->skip_num_rows(offset_);
         }
-        block->swap(sorted_blocks_[0]);
+        block->swap(*priority_queue_.top().impl->block);
         *eos = true;
     } else {
         RETURN_IF_ERROR(_merge_sort_read_impl(batch_size, block, eos));
@@ -110,9 +106,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'.
@@ -123,7 +124,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_--;
@@ -134,7 +135,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));
 
@@ -261,22 +264,22 @@ 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)));
+            _state->add_sorted_block(Block::create_shared(std::move(desc_block)));
+            _block_priority_queue.emplace(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());
+            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(desc_block));
-                _block_priority_queue.emplace(_pool->add(
-                        new MergeSortCursorImpl(_state->last_sorted_block(), _sort_description)));
+                _state->add_sorted_block(tmp_cursor_impl->block);
+                _block_priority_queue.emplace(MergeSortCursorImpl::create_shared(
+                        _state->last_sorted_block(), _sort_description));
             }
         }
     } else {
         // dispose normal sort logic
-        static_cast<void>(_state->add_sorted_block(desc_block));
+        _state->add_sorted_block(Block::create_shared(std::move(desc_block)));
     }
     return Status::OK();
 }
diff --git a/be/src/vec/common/sort/sorter.h b/be/src/vec/common/sort/sorter.h
@@ -59,7 +59,7 @@ class MergeSorterState {
 
     ~MergeSorterState() = default;
 
-    Status add_sorted_block(Block& block);
+    void add_sorted_block(std::shared_ptr<Block> block);
 
     Status build_merge_tree(const SortDescription& sort_description);
 
@@ -72,23 +72,19 @@ 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_; }
     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<Block>> sorted_blocks_;
     size_t in_mem_sorted_bocks_size_ = 0;
     uint64_t num_rows_ = 0;
 
diff --git a/be/src/vec/common/sort/topn_sorter.cpp b/be/src/vec/common/sort/topn_sorter.cpp
@@ -72,17 +72,16 @@ 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)));
+            _state->add_sorted_block(Block::create_shared(std::move(sorted_block)));
+            _block_priority_queue.emplace(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());
+            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(sorted_block));
-                _block_priority_queue.emplace(_pool->add(
-                        new MergeSortCursorImpl(_state->last_sorted_block(), _sort_description)));
+                _state->add_sorted_block(block_cursor.impl->block);
+                _block_priority_queue.emplace(tmp_cursor_impl);
             }
         }
     } else {
diff --git a/be/src/vec/core/sort_cursor.h b/be/src/vec/core/sort_cursor.h
diff --git a/be/src/vec/runtime/vsorted_run_merger.cpp b/be/src/vec/runtime/vsorted_run_merger.cpp
diff --git a/be/src/vec/runtime/vsorted_run_merger.h b/be/src/vec/runtime/vsorted_run_merger.h
