branch-4.1: [fix](be) Continue sorted merge when sender queue is ready #65004

be/src/exec/exchange/vdata_stream_recvr.cpp

@@ -130,6 +130,11 @@ Status VDataStreamRecvr::SenderQueue::get_batch(Block* block, bool* eos) {
     return Status::OK();
 }
 
+bool VDataStreamRecvr::SenderQueue::has_data_or_finished() {
+    std::lock_guard<std::mutex> l(_lock);
+    return _is_cancelled || !_block_queue.empty() || _num_remaining_senders == 0;
+}
+
 void VDataStreamRecvr::SenderQueue::set_source_ready(std::lock_guard<std::mutex>&) {
     // Here, it is necessary to check if _source_dependency is not nullptr.
     // This is because the queue might be closed before setting the source dependency.
@@ -391,16 +396,21 @@ Status VDataStreamRecvr::create_merger(const VExprContextSPtrs& ordering_expr,
     DCHECK(_is_merging);
     SCOPED_CONSUME_MEM_TRACKER(_mem_tracker.get());
     std::vector<BlockSupplier> child_block_suppliers;
+    std::vector<BlockSupplierReadyChecker> child_block_supplier_ready_checkers;
     // Create the merger that will a single stream of sorted rows.
     _merger.reset(new VSortedRunMerger(ordering_expr, is_asc_order, nulls_first, batch_size, limit,
                                        offset, _profile));
 
+    child_block_suppliers.reserve(_sender_queues.size());
+    child_block_supplier_ready_checkers.reserve(_sender_queues.size());
     for (int i = 0; i < _sender_queues.size(); ++i) {
         child_block_suppliers.emplace_back(std::bind(std::mem_fn(&SenderQueue::get_batch),
                                                      _sender_queues[i], std::placeholders::_1,
                                                      std::placeholders::_2));
+        child_block_supplier_ready_checkers.emplace_back(
+                std::bind(std::mem_fn(&SenderQueue::has_data_or_finished), _sender_queues[i]));
     }
-    RETURN_IF_ERROR(_merger->prepare(child_block_suppliers));
+    RETURN_IF_ERROR(_merger->prepare(child_block_suppliers, child_block_supplier_ready_checkers));
     return Status::OK();
 }
 

be/src/exec/exchange/vdata_stream_recvr.h

@@ -179,6 +179,8 @@ class VDataStreamRecvr::SenderQueue {
 
     Status get_batch(Block* next_block, bool* eos);
 
+    bool has_data_or_finished();
+
     Status add_block(std::unique_ptr<PBlock> pblock, int be_number, int64_t packet_seq,
                      ::google::protobuf::Closure** done, const int64_t wait_for_worker,
                      const uint64_t time_to_find_recvr);

be/src/exec/sort/sort_cursor.h

@@ -22,6 +22,7 @@
 
 #include <glog/logging.h>
 
+#include <functional>
 #include <utility>
 
 #include "core/block/block.h"
@@ -108,6 +109,7 @@ struct MergeSortCursorImpl {
     virtual void process_next() {}
     virtual Block* block_ptr() { return nullptr; }
     virtual bool eof() const { return false; }
+    virtual bool has_ready_block_or_eos() const { return false; }
 
     Field get_top_value() const {
         Field field {PrimitiveType::TYPE_NULL};
@@ -117,14 +119,18 @@ struct MergeSortCursorImpl {
 };
 
 using BlockSupplier = std::function<Status(Block*, bool* eos)>;
+using BlockSupplierReadyChecker = std::function<bool()>;
 
 struct BlockSupplierSortCursorImpl : public MergeSortCursorImpl {
     ENABLE_FACTORY_CREATOR(BlockSupplierSortCursorImpl);
     BlockSupplierSortCursorImpl(BlockSupplier block_supplier,
                                 const VExprContextSPtrs& ordering_expr,
                                 const std::vector<bool>& is_asc_order,
-                                const std::vector<bool>& nulls_first)
-            : _ordering_expr(ordering_expr), _block_supplier(std::move(block_supplier)) {
+                                const std::vector<bool>& nulls_first,
+                                BlockSupplierReadyChecker block_supplier_ready_checker = {})
+            : _ordering_expr(ordering_expr),
+              _block_supplier(std::move(block_supplier)),
+              _block_supplier_ready_checker(std::move(block_supplier_ready_checker)) {
         block = Block::create_shared();
         sort_columns_size = ordering_expr.size();
 
@@ -136,11 +142,18 @@ struct BlockSupplierSortCursorImpl : public MergeSortCursorImpl {
         process_next();
     }
 
-    BlockSupplierSortCursorImpl(BlockSupplier block_supplier, const SortDescription& desc_)
-            : MergeSortCursorImpl(desc_), _block_supplier(std::move(block_supplier)) {
+    BlockSupplierSortCursorImpl(BlockSupplier block_supplier, const SortDescription& desc_,
+                                BlockSupplierReadyChecker block_supplier_ready_checker = {})
+            : MergeSortCursorImpl(desc_),
+              _block_supplier(std::move(block_supplier)),
+              _block_supplier_ready_checker(std::move(block_supplier_ready_checker)) {
         process_next();
     }
 
+    bool has_ready_block_or_eos() const override {
+        return _block_supplier_ready_checker && _block_supplier_ready_checker();
+    }
+
     void process_next() override {
         if (_is_eof) {
             return;
@@ -167,6 +180,7 @@ struct BlockSupplierSortCursorImpl : public MergeSortCursorImpl {
 
     VExprContextSPtrs _ordering_expr;
     BlockSupplier _block_supplier;
+    BlockSupplierReadyChecker _block_supplier_ready_checker;
     bool _is_eof = false;
 };
 

be/src/exec/sort/vsorted_run_merger.cpp

@@ -62,14 +62,22 @@ void VSortedRunMerger::init_timers(RuntimeProfile* profile) {
     _get_next_block_timer = ADD_TIMER(profile, "MergeGetNextBlock");
 }
 
-Status VSortedRunMerger::prepare(const std::vector<BlockSupplier>& input_runs) {
+Status VSortedRunMerger::prepare(
+        const std::vector<BlockSupplier>& input_runs,
+        const std::vector<BlockSupplierReadyChecker>& input_run_ready_checkers) {
     try {
-        for (const auto& supplier : input_runs) {
+        for (size_t i = 0; i < input_runs.size(); ++i) {
+            const auto& supplier = input_runs[i];
+            BlockSupplierReadyChecker ready_checker;
+            if (i < input_run_ready_checkers.size()) {
+                ready_checker = input_run_ready_checkers[i];
+            }
             if (_use_sort_desc) {
-                _cursors.emplace_back(BlockSupplierSortCursorImpl::create_shared(supplier, _desc));
+                _cursors.emplace_back(
+                        BlockSupplierSortCursorImpl::create_shared(supplier, _desc, ready_checker));
             } else {
                 _cursors.emplace_back(BlockSupplierSortCursorImpl::create_shared(
-                        supplier, _ordering_expr, _is_asc_order, _nulls_first));
+                        supplier, _ordering_expr, _is_asc_order, _nulls_first, ready_checker));
             }
         }
     } catch (const std::exception& e) {
@@ -194,7 +202,12 @@ Status VSortedRunMerger::get_next(Block* output_block, bool* eos) {
             }
 
             current->next();
-            if (_need_more_data(current)) {
+            const bool has_ready_block_or_eos =
+                    current->is_last(0) && !current->eof() && current->has_ready_block_or_eos();
+            if (has_ready_block_or_eos) {
+                do_insert();
+            }
+            if (_need_more_data(current, has_ready_block_or_eos)) {
                 do_insert();
                 scoped_mutable_block.restore();
                 return Status::OK();
@@ -212,12 +225,21 @@ Status VSortedRunMerger::get_next(Block* output_block, bool* eos) {
     return Status::OK();
 }
 
-bool VSortedRunMerger::_need_more_data(MergeSortCursor& current) {
+bool VSortedRunMerger::_need_more_data(MergeSortCursor& current, bool has_ready_block_or_eos) {
     if (!current->is_last(0)) {
         _priority_queue.push(current);
         return false;
     } else if (current->eof()) {
         return false;
+    } else if (has_ready_block_or_eos) {
+        {
+            ScopedTimer<MonotonicStopWatch> timer(_get_next_block_timer);
+            current->process_next();
+        }
+        if (!current->eof()) {
+            _priority_queue.push(current);
+        }
+        return false;
     } else {
         _pending_cursor = current.impl;
         return true;

be/src/exec/sort/vsorted_run_merger.h

@@ -56,7 +56,8 @@ class VSortedRunMerger {
     // Prepare this merger to merge and return rows from the sorted runs in 'input_runs'.
     // Retrieves the first batch from each run and sets up the binary heap implementing
     // the priority queue.
-    Status prepare(const std::vector<BlockSupplier>& input_runs);
+    Status prepare(const std::vector<BlockSupplier>& input_runs,
+                   const std::vector<BlockSupplierReadyChecker>& input_run_ready_checkers = {});
 
     // Return the next block of sorted rows from this merger.
     Status get_next(Block* output_block, bool* eos);
@@ -92,8 +93,9 @@ class VSortedRunMerger {
 
 private:
     void init_timers(RuntimeProfile* profile);
-    // If current stream is exhausted and not eof, we should break this loop and read more blocks.
-    bool _need_more_data(MergeSortCursor& current);
+    // If current stream is exhausted and not eof, keep merging only when the next block or eos is
+    // already available. Otherwise break this loop and wait for more data.
+    bool _need_more_data(MergeSortCursor& current, bool has_ready_block_or_eos);
 };
 
 } // namespace doris