diff --git a/be/src/exec/exchange/vdata_stream_recvr.cpp b/be/src/exec/exchange/vdata_stream_recvr.cpp index d98b49a9e09500..ebce52683cd34d 100644 --- a/be/src/exec/exchange/vdata_stream_recvr.cpp +++ b/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 l(_lock); + return _is_cancelled || !_block_queue.empty() || _num_remaining_senders == 0; +} + void VDataStreamRecvr::SenderQueue::set_source_ready(std::lock_guard&) { // 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 child_block_suppliers; + std::vector 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(); } diff --git a/be/src/exec/exchange/vdata_stream_recvr.h b/be/src/exec/exchange/vdata_stream_recvr.h index 88e2604a601d92..c465e8b597cbde 100644 --- a/be/src/exec/exchange/vdata_stream_recvr.h +++ b/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, int be_number, int64_t packet_seq, ::google::protobuf::Closure** done, const int64_t wait_for_worker, const uint64_t time_to_find_recvr); diff --git a/be/src/exec/sort/sort_cursor.h b/be/src/exec/sort/sort_cursor.h index d5b4a14e46158f..6d30c4a935110f 100644 --- a/be/src/exec/sort/sort_cursor.h +++ b/be/src/exec/sort/sort_cursor.h @@ -22,6 +22,7 @@ #include +#include #include #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; +using BlockSupplierReadyChecker = std::function; struct BlockSupplierSortCursorImpl : public MergeSortCursorImpl { ENABLE_FACTORY_CREATOR(BlockSupplierSortCursorImpl); BlockSupplierSortCursorImpl(BlockSupplier block_supplier, const VExprContextSPtrs& ordering_expr, const std::vector& is_asc_order, - const std::vector& nulls_first) - : _ordering_expr(ordering_expr), _block_supplier(std::move(block_supplier)) { + const std::vector& 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; }; diff --git a/be/src/exec/sort/vsorted_run_merger.cpp b/be/src/exec/sort/vsorted_run_merger.cpp index 4981b0ae3ea878..928bb99c16b335 100644 --- a/be/src/exec/sort/vsorted_run_merger.cpp +++ b/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& input_runs) { +Status VSortedRunMerger::prepare( + const std::vector& input_runs, + const std::vector& 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 timer(_get_next_block_timer); + current->process_next(); + } + if (!current->eof()) { + _priority_queue.push(current); + } + return false; } else { _pending_cursor = current.impl; return true; diff --git a/be/src/exec/sort/vsorted_run_merger.h b/be/src/exec/sort/vsorted_run_merger.h index 471333e862e5c3..d06510f35aa700 100644 --- a/be/src/exec/sort/vsorted_run_merger.h +++ b/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& input_runs); + Status prepare(const std::vector& input_runs, + const std::vector& 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 diff --git a/be/test/core/value/sort_merger_test.cpp b/be/test/core/value/sort_merger_test.cpp index b5c0b8ced95e2a..75a95db7c9b291 100644 --- a/be/test/core/value/sort_merger_test.cpp +++ b/be/test/core/value/sort_merger_test.cpp @@ -17,8 +17,13 @@ #include +#include +#include +#include + #include "core/data_type/data_type_nullable.h" #include "core/data_type/data_type_number.h" +#include "core/data_type/data_type_string.h" #include "exec/sort/vsorted_run_merger.h" #include "testutil/column_helper.h" #include "testutil/mock/mock_slot_ref.h" @@ -33,6 +38,16 @@ class SortMergerTest : public testing::Test { void TearDown() override {} }; +static Block create_int_string_block(const std::vector& keys, + const std::vector& values) { + DCHECK_EQ(keys.size(), values.size()); + auto key_column = ColumnHelper::create_column(keys); + auto value_column = ColumnHelper::create_column(values); + Block block({ColumnWithTypeAndName(key_column, std::make_shared(), "key"), + ColumnWithTypeAndName(value_column, std::make_shared(), "value")}); + return block; +} + TEST(SortMergerTest, NULL_FIRST_ASC) { /** * in: [([NULL, 1, 2, 3, 4], eos = false), ([], eos = true)] @@ -478,4 +493,188 @@ TEST(SortMergerTest, TEST_SINGLE_STREAM) { } } -} // namespace doris \ No newline at end of file +TEST(SortMergerTest, CONTINUE_MERGE_WHEN_NEXT_BLOCK_READY) { + const int batch_size = 4; + std::vector>> input_blocks = {{{1}, {3}}, {{2}, {4}}}; + std::vector next_block_index(input_blocks.size(), 0); + + std::unique_ptr merger; + auto profile = std::make_shared(""); + auto ordering_expr = MockSlotRef::create_mock_contexts(std::make_shared()); + { + std::vector is_asc_order = {true}; + std::vector nulls_first = {true}; + const int limit = -1; + const int offset = 0; + merger.reset(new VSortedRunMerger(ordering_expr, is_asc_order, nulls_first, batch_size, + limit, offset, profile.get())); + } + { + std::vector child_block_suppliers; + std::vector ready_checkers; + for (size_t child_idx = 0; child_idx < input_blocks.size(); child_idx++) { + BlockSupplier block_supplier = [&, id = child_idx](Block* block, bool* eos) { + if (next_block_index[id] >= input_blocks[id].size()) { + *eos = true; + return Status::OK(); + } + *block = ColumnHelper::create_block( + input_blocks[id][next_block_index[id]++]); + *eos = false; + return Status::OK(); + }; + child_block_suppliers.push_back(block_supplier); + ready_checkers.emplace_back([] { return true; }); + } + EXPECT_TRUE(merger->prepare(child_block_suppliers, ready_checkers).ok()); + } + { + Block block; + bool eos = false; + EXPECT_TRUE(merger->get_next(&block, &eos).ok()); + auto expect_block = ColumnHelper::create_column({1, 2, 3, 4}); + EXPECT_TRUE(ColumnHelper::column_equal(block.get_by_position(0).column, expect_block)); + EXPECT_EQ(block.rows(), batch_size); + EXPECT_FALSE(eos); + } + { + Block block; + bool eos = false; + EXPECT_TRUE(merger->get_next(&block, &eos).ok()); + EXPECT_EQ(block.rows(), 0); + EXPECT_TRUE(eos); + } +} + +TEST(SortMergerTest, CONTINUE_MERGE_WITH_STRING_COLUMN_WHEN_NEXT_BLOCK_READY) { + const int batch_size = 4; + std::vector> input_blocks; + input_blocks.emplace_back( + std::vector {create_int_string_block({1, 3}, {"old-1", "old-3"}), + create_int_string_block({5, 7}, {"new-5", "new-7"})}); + input_blocks.emplace_back( + std::vector {create_int_string_block({2, 4}, {"other-2", "other-4"})}); + std::vector next_block_index(input_blocks.size(), 0); + + std::unique_ptr merger; + auto profile = std::make_shared(""); + auto ordering_expr = MockSlotRef::create_mock_contexts(std::make_shared()); + { + std::vector is_asc_order = {true}; + std::vector nulls_first = {true}; + const int limit = -1; + const int offset = 0; + merger.reset(new VSortedRunMerger(ordering_expr, is_asc_order, nulls_first, batch_size, + limit, offset, profile.get())); + } + { + std::vector child_block_suppliers; + std::vector ready_checkers; + for (size_t child_idx = 0; child_idx < input_blocks.size(); child_idx++) { + BlockSupplier block_supplier = [&, id = child_idx](Block* block, bool* eos) { + if (next_block_index[id] >= input_blocks[id].size()) { + *eos = true; + return Status::OK(); + } + block->swap(input_blocks[id][next_block_index[id]++]); + *eos = false; + return Status::OK(); + }; + child_block_suppliers.push_back(block_supplier); + ready_checkers.emplace_back([] { return true; }); + } + EXPECT_TRUE(merger->prepare(child_block_suppliers, ready_checkers).ok()); + } + { + Block block; + bool eos = false; + EXPECT_TRUE(merger->get_next(&block, &eos).ok()); + auto expect_key_column = ColumnHelper::create_column({1, 2, 3, 4}); + auto expect_value_column = ColumnHelper::create_column( + {"old-1", "other-2", "old-3", "other-4"}); + EXPECT_TRUE(ColumnHelper::column_equal(block.get_by_position(0).column, expect_key_column)); + EXPECT_TRUE( + ColumnHelper::column_equal(block.get_by_position(1).column, expect_value_column)); + EXPECT_EQ(block.rows(), batch_size); + EXPECT_FALSE(eos); + } + { + Block block; + bool eos = false; + EXPECT_TRUE(merger->get_next(&block, &eos).ok()); + auto expect_key_column = ColumnHelper::create_column({5, 7}); + auto expect_value_column = ColumnHelper::create_column({"new-5", "new-7"}); + EXPECT_TRUE(ColumnHelper::column_equal(block.get_by_position(0).column, expect_key_column)); + EXPECT_TRUE( + ColumnHelper::column_equal(block.get_by_position(1).column, expect_value_column)); + EXPECT_EQ(block.rows(), 2); + EXPECT_FALSE(eos); + } +} + +TEST(SortMergerTest, KEEP_PENDING_CURSOR_WHEN_NEXT_BLOCK_NOT_READY) { + const int batch_size = 4; + std::vector>> input_blocks = {{{1}, {3}}, {{2}, {4}}}; + std::vector next_block_index(input_blocks.size(), 0); + std::vector ready(input_blocks.size(), false); + + std::unique_ptr merger; + auto profile = std::make_shared(""); + auto ordering_expr = MockSlotRef::create_mock_contexts(std::make_shared()); + { + std::vector is_asc_order = {true}; + std::vector nulls_first = {true}; + const int limit = -1; + const int offset = 0; + merger.reset(new VSortedRunMerger(ordering_expr, is_asc_order, nulls_first, batch_size, + limit, offset, profile.get())); + } + { + std::vector child_block_suppliers; + std::vector ready_checkers; + for (size_t child_idx = 0; child_idx < input_blocks.size(); child_idx++) { + BlockSupplier block_supplier = [&, id = child_idx](Block* block, bool* eos) { + if (next_block_index[id] >= input_blocks[id].size()) { + *eos = true; + return Status::OK(); + } + *block = ColumnHelper::create_block( + input_blocks[id][next_block_index[id]++]); + *eos = false; + return Status::OK(); + }; + child_block_suppliers.push_back(block_supplier); + ready_checkers.emplace_back([&, id = child_idx] { return ready[id]; }); + } + EXPECT_TRUE(merger->prepare(child_block_suppliers, ready_checkers).ok()); + } + { + Block block; + bool eos = false; + EXPECT_TRUE(merger->get_next(&block, &eos).ok()); + auto expect_block = ColumnHelper::create_column({1}); + EXPECT_TRUE(ColumnHelper::column_equal(block.get_by_position(0).column, expect_block)); + EXPECT_EQ(block.rows(), 1); + EXPECT_FALSE(eos); + } + { + ready[0] = true; + ready[1] = true; + Block block; + bool eos = false; + EXPECT_TRUE(merger->get_next(&block, &eos).ok()); + auto expect_block = ColumnHelper::create_column({2, 3, 4}); + EXPECT_TRUE(ColumnHelper::column_equal(block.get_by_position(0).column, expect_block)); + EXPECT_EQ(block.rows(), 3); + EXPECT_FALSE(eos); + } + { + Block block; + bool eos = false; + EXPECT_TRUE(merger->get_next(&block, &eos).ok()); + EXPECT_EQ(block.rows(), 0); + EXPECT_TRUE(eos); + } +} + +} // namespace doris