This is an automated email from the ASF dual-hosted git repository.
Mryange 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 039a1f0909f [refine](exec) simplify DataQueue eos handling (#65324)
039a1f0909f is described below
commit 039a1f0909f14a2c1942e9605fe904e2b22d34e9
Author: Mryange <[email protected]>
AuthorDate: Fri Jul 17 14:41:11 2026 +0800
[refine](exec) simplify DataQueue eos handling (#65324)
### What problem does this PR solve?
DataQueue exposed its internal child-finish and queue-selection protocol
to source and sink operators, which made EOS handling duplicated and
difficult to follow.
This change carries EOS through `push_block()` and `DataQueueBlock`,
moves child selection into DataQueue, and updates Union and Cache
operators to consume the unified result. Union sink also forwards each
input block immediately instead of retaining a cross-call output
### Release note
None
---
be/src/exec/operator/cache_sink_operator.cpp | 13 ++-
be/src/exec/operator/cache_source_operator.cpp | 17 +--
be/src/exec/operator/cache_source_operator.h | 4 -
be/src/exec/operator/data_queue.cpp | 103 ++++++++++-------
be/src/exec/operator/data_queue.h | 31 ++++--
be/src/exec/operator/union_sink_operator.cpp | 37 ++-----
be/src/exec/operator/union_sink_operator.h | 8 +-
be/src/exec/operator/union_source_operator.cpp | 41 ++-----
be/src/exec/operator/union_source_operator.h | 9 +-
be/test/exec/operator/union_operator_test.cpp | 6 +-
be/test/exec/pipeline/data_queue_test.cpp | 148 ++++++++++++++-----------
11 files changed, 201 insertions(+), 216 deletions(-)
diff --git a/be/src/exec/operator/cache_sink_operator.cpp
b/be/src/exec/operator/cache_sink_operator.cpp
index 97c42b99c25..b8c7204ec1a 100644
--- a/be/src/exec/operator/cache_sink_operator.cpp
+++ b/be/src/exec/operator/cache_sink_operator.cpp
@@ -54,12 +54,13 @@ Status CacheSinkOperatorX::sink_impl(RuntimeState* state,
Block* in_block, bool
SCOPED_TIMER(local_state.exec_time_counter());
COUNTER_UPDATE(local_state.rows_input_counter(),
(int64_t)in_block->rows());
- if (in_block->rows() > 0) {
- RETURN_IF_ERROR(local_state._shared_state->data_queue.push_block(
- Block::create_unique(std::move(*in_block)), 0));
- }
- if (UNLIKELY(eos)) {
- local_state._shared_state->data_queue.set_finish(0);
+ if (in_block->rows() > 0 || eos) {
+ std::unique_ptr<Block> output_block;
+ if (in_block->rows() > 0) {
+ output_block = Block::create_unique(std::move(*in_block));
+ }
+ RETURN_IF_ERROR(
+
local_state._shared_state->data_queue.push_block(std::move(output_block), 0,
eos));
}
return Status::OK();
}
diff --git a/be/src/exec/operator/cache_source_operator.cpp
b/be/src/exec/operator/cache_source_operator.cpp
index ec7d947680a..6b68e300ff9 100644
--- a/be/src/exec/operator/cache_source_operator.cpp
+++ b/be/src/exec/operator/cache_source_operator.cpp
@@ -111,9 +111,8 @@ std::string CacheSourceLocalState::debug_string(int
indentation_level) const {
fmt::memory_buffer debug_string_buffer;
fmt::format_to(debug_string_buffer, "{}",
Base::debug_string(indentation_level));
if (_shared_state) {
- fmt::format_to(debug_string_buffer, ", data_queue: (is_all_finish =
{}, has_data = {})",
- _shared_state->data_queue.is_all_finish(),
- _shared_state->data_queue.has_more_data());
+ fmt::format_to(debug_string_buffer, ", data_queue: {}",
+ _shared_state->data_queue.debug_string());
}
return fmt::to_string(debug_string_buffer);
}
@@ -140,18 +139,14 @@ Status CacheSourceOperatorX::get_block_impl(RuntimeState*
state, Block* block, b
}
});
- std::unique_ptr<Block> output_block;
- int child_idx = 0;
-
RETURN_IF_ERROR(local_state._shared_state->data_queue.get_block_from_queue(&output_block,
-
&child_idx));
- // Here, check the value of `_has_data(state)` again after
`data_queue.is_all_finish()` is TRUE
- // as there may be one or more blocks when
`data_queue.is_all_finish()` is TRUE.
- *eos = !_has_data(state) &&
local_state._shared_state->data_queue.is_all_finish();
+ auto queue_block =
DORIS_TRY(local_state._shared_state->data_queue.get_block_from_queue());
+ *eos = queue_block.eos;
- if (!output_block) {
+ if (!queue_block.block) {
return Status::OK();
}
+ auto& output_block = queue_block.block;
if (local_state._need_insert_cache) {
if (need_clone_empty) {
*block = output_block->clone_empty();
diff --git a/be/src/exec/operator/cache_source_operator.h
b/be/src/exec/operator/cache_source_operator.h
index 3a68bd337fc..864f6301217 100644
--- a/be/src/exec/operator/cache_source_operator.h
+++ b/be/src/exec/operator/cache_source_operator.h
@@ -90,10 +90,6 @@ public:
private:
TQueryCacheParam _cache_param;
- bool _has_data(RuntimeState* state) const {
- auto& local_state = get_local_state(state);
- return local_state._shared_state->data_queue.remaining_has_data();
- }
friend class CacheSourceLocalState;
};
diff --git a/be/src/exec/operator/data_queue.cpp
b/be/src/exec/operator/data_queue.cpp
index 752e972c3e6..feed0bb79fa 100644
--- a/be/src/exec/operator/data_queue.cpp
+++ b/be/src/exec/operator/data_queue.cpp
@@ -25,6 +25,7 @@
#include "common/thread_safety_annotations.h"
#include "core/block/block.h"
#include "exec/pipeline/dependency.h"
+#include "fmt/format.h"
namespace doris {
@@ -97,6 +98,10 @@ bool DataQueue::has_more_data() const {
return _cur_blocks_total_nums.load() > 0;
}
+std::string DataQueue::debug_string() const {
+ return fmt::format("(is_all_finish = {}, has_data = {})", is_all_finish(),
has_more_data());
+}
+
void DataQueue::set_source_dependency(std::shared_ptr<Dependency>
source_dependency)
NO_THREAD_SAFETY_ANALYSIS {
_source_dependency = std::move(source_dependency);
@@ -134,13 +139,16 @@ std::unique_ptr<Block> DataQueue::get_free_block(int
child_idx) {
return Block::create_unique();
}
-void DataQueue::push_free_block(std::unique_ptr<Block> block, int child_idx) {
- DCHECK(block->rows() == 0);
+void DataQueue::push_free_block(DataQueueBlock&& queue_block) {
+ if (!queue_block.block) {
+ return;
+ }
+ DCHECK(queue_block.block->rows() == 0);
if (!_is_low_memory_mode) {
- auto& sub = *_sub_queues[child_idx];
+ auto& sub = *_sub_queues[queue_block.child_idx];
LockGuard l(sub.free_lock);
- sub.free_blocks.emplace_back(std::move(block));
+ sub.free_blocks.emplace_back(std::move(queue_block.block));
}
}
@@ -154,73 +162,75 @@ void DataQueue::clear_free_blocks() {
void DataQueue::terminate() {
for (int i = 0; i < _child_count; ++i) {
- set_finish(i);
+ mark_finish(i);
_sub_queues[i]->clear_blocks();
}
+ _cur_blocks_total_nums = 0;
clear_free_blocks();
+ set_source_always_ready();
}
-//check which queue have data, and save the idx in _flag_queue_idx,
-//so next loop, will check the record idx + 1 first
-//maybe it's useful with many queue, others maybe always 0
-bool DataQueue::remaining_has_data() {
- int count = _child_count;
- while (--count >= 0) {
- _flag_queue_idx++;
- if (_flag_queue_idx == _child_count) {
- _flag_queue_idx = 0;
- }
- if (_sub_queues[_flag_queue_idx]->blocks_in_queue.load() > 0) {
- return true;
+Result<DataQueueBlock> DataQueue::get_block_from_queue() {
+ DataQueueBlock result;
+ const int start_idx = (_flag_queue_idx + 1) % _child_count;
+ for (int offset = 0; offset < _child_count; ++offset) {
+ const int idx = (start_idx + offset) % _child_count;
+ if (_sub_queues[idx]->blocks_in_queue.load() == 0) {
+ continue;
}
- }
- return false;
-}
-
-//the _flag_queue_idx indicate which queue has data, and in check can_read
-//will be set idx in remaining_has_data function
-Status DataQueue::get_block_from_queue(std::unique_ptr<Block>* output_block,
int* child_idx) {
- const int idx = _flag_queue_idx;
- auto& sub = *_sub_queues[idx];
- sub.try_pop(output_block);
- if (*output_block) {
- if (child_idx) {
- *child_idx = idx;
+ auto& sub = *_sub_queues[idx];
+ sub.try_pop(&result.block);
+ if (!result.block) {
+ continue;
}
+ result.child_idx = idx;
+ _flag_queue_idx = idx;
auto old_total = _cur_blocks_total_nums.fetch_sub(1);
if (old_total == 1) {
set_source_block();
}
+ break;
}
- return Status::OK();
+
+ // A producer enqueues its final block before marking the child finished.
Observe completion
+ // first, then recheck queued data to avoid reporting EOS while the final
block is still queued.
+ result.eos = is_all_finish() && !has_more_data();
+ return result;
}
-Status DataQueue::push_block(std::unique_ptr<Block> block, int child_idx) {
- if (!block) {
+Status DataQueue::push_block(std::unique_ptr<Block> block, int child_idx, bool
eos) {
+ DCHECK(block || eos);
+ if (!block && !eos) {
return Status::OK();
}
- auto& sub = *_sub_queues[child_idx];
- // total_counter is incremented inside try_push under queue_lock, only
when the
- // block is actually enqueued. This ensures get_block_from_queue() always
observes
- // _cur_blocks_total_nums >= 1 when it successfully pops a block, with no
risk of
- // underflow or the need for a rollback on failure.
- if (!sub.try_push(std::move(block), _cur_blocks_total_nums)) {
- return Status::EndOfFile("SubQueue already finished");
+
+ if (block) {
+ auto& sub = *_sub_queues[child_idx];
+ // total_counter is incremented inside try_push under queue_lock, only
when the
+ // block is actually enqueued. This ensures get_block_from_queue()
always observes
+ // _cur_blocks_total_nums >= 1 when it successfully pops a block, with
no risk of
+ // underflow or the need for a rollback on failure.
+ if (!sub.try_push(std::move(block), _cur_blocks_total_nums)) {
+ return Status::EndOfFile("SubQueue already finished");
+ }
+ }
+
+ if (eos) {
+ mark_finish(child_idx);
}
set_source_ready();
return Status::OK();
}
-void DataQueue::set_finish(int child_idx) {
+void DataQueue::mark_finish(int child_idx) {
auto& sub = *_sub_queues[child_idx];
if (!sub.mark_finished(_un_finished_counter, _is_all_finished)) {
return;
}
- set_source_ready();
}
-bool DataQueue::is_all_finish() {
+bool DataQueue::is_all_finish() const {
return _is_all_finished;
}
@@ -231,6 +241,13 @@ void DataQueue::set_source_ready() {
}
}
+void DataQueue::set_source_always_ready() {
+ LockGuard lc(_source_lock);
+ if (_source_dependency) {
+ _source_dependency->set_always_ready();
+ }
+}
+
void DataQueue::set_source_block() {
// Re-check under _source_lock to avoid blocking the source when a
concurrent push
// has already added new blocks (or all children have finished) since we
observed
diff --git a/be/src/exec/operator/data_queue.h
b/be/src/exec/operator/data_queue.h
index 1443e2ed520..49a3bcaf00d 100644
--- a/be/src/exec/operator/data_queue.h
+++ b/be/src/exec/operator/data_queue.h
@@ -20,6 +20,7 @@
#include <cstdint>
#include <deque>
#include <memory>
+#include <string>
#include <vector>
#include "common/status.h"
@@ -30,6 +31,15 @@ namespace doris {
class Dependency;
+struct DataQueueBlock {
+ std::unique_ptr<Block> block;
+ bool eos = false;
+
+private:
+ int child_idx = 0;
+ friend class DataQueue;
+};
+
// Per child sub-queue. Groups all parallel state so that the lock/field
// relationship is explicit and can be checked by clang -Wthread-safety.
struct SubQueue {
@@ -46,7 +56,7 @@ struct SubQueue {
AnnotatedMutex free_lock;
std::deque<std::unique_ptr<Block>> free_blocks GUARDED_BY(free_lock);
- // blocks_in_queue is readable from lock-free fast paths
(remaining_has_data),
+ // blocks_in_queue is readable from lock-free fast paths
(get_block_from_queue),
// so it remains atomic and is intentionally not GUARDED_BY.
std::atomic_uint32_t blocks_in_queue {0};
@@ -82,18 +92,13 @@ public:
DataQueue(int child_count = 1);
~DataQueue() = default;
- Status get_block_from_queue(std::unique_ptr<Block>* block, int* child_idx
= nullptr);
- Status push_block(std::unique_ptr<Block> block, int child_idx = 0);
+ Result<DataQueueBlock> get_block_from_queue();
+ Status push_block(std::unique_ptr<Block> block, int child_idx, bool eos);
std::unique_ptr<Block> get_free_block(int child_idx = 0);
- void push_free_block(std::unique_ptr<Block> output_block, int child_idx =
0);
+ void push_free_block(DataQueueBlock&& queue_block);
- void set_finish(int child_idx = 0);
- bool is_all_finish();
-
- // This function is not thread safe, should be called in
Operator::get_block()
- bool remaining_has_data();
- bool has_more_data() const;
+ std::string debug_string() const;
void set_source_dependency(std::shared_ptr<Dependency> source_dependency)
NO_THREAD_SAFETY_ANALYSIS;
@@ -104,8 +109,12 @@ public:
void terminate();
private:
+ bool is_all_finish() const;
+ bool has_more_data() const;
void clear_free_blocks();
+ void mark_finish(int child_idx);
void set_source_ready();
+ void set_source_always_ready();
void set_source_block();
std::vector<std::unique_ptr<SubQueue>> _sub_queues;
@@ -117,7 +126,7 @@ private:
std::atomic_uint32_t _cur_blocks_total_nums = 0;
//this will be indicate which queue has data, it's useful when have many
queues
- std::atomic_int _flag_queue_idx = 0;
+ int _flag_queue_idx = 0;
// only used by streaming agg source operator
std::atomic_bool _is_low_memory_mode = false;
diff --git a/be/src/exec/operator/union_sink_operator.cpp
b/be/src/exec/operator/union_sink_operator.cpp
index b0f7cf619a8..81881fcf576 100644
--- a/be/src/exec/operator/union_sink_operator.cpp
+++ b/be/src/exec/operator/union_sink_operator.cpp
@@ -100,44 +100,21 @@ Status UnionSinkOperatorX::sink_impl(RuntimeState* state,
Block* in_block, bool
}
SCOPED_TIMER(local_state.exec_time_counter());
COUNTER_UPDATE(local_state.rows_input_counter(),
(int64_t)in_block->rows());
- if (local_state._output_block == nullptr) {
- local_state._output_block =
-
local_state._shared_state->data_queue.get_free_block(_cur_child_id);
- }
- if (_cur_child_id < _get_first_materialized_child_idx()) { //pass_through
- if (in_block->rows() > 0) {
- local_state._output_block->swap(*in_block);
- RETURN_IF_ERROR(local_state._shared_state->data_queue.push_block(
- std::move(local_state._output_block), _cur_child_id));
- }
+ auto output_block =
local_state._shared_state->data_queue.get_free_block(_cur_child_id);
+ if (is_child_passthrough(_cur_child_id)) { //pass_through
+ output_block->swap(*in_block);
} else if (_get_first_materialized_child_idx() != children_count() &&
_cur_child_id < children_count()) { //need materialized
- RETURN_IF_ERROR(materialize_child_block(state, _cur_child_id, in_block,
-
local_state._output_block.get()));
+ RETURN_IF_ERROR(
+ materialize_child_block(state, _cur_child_id, in_block,
output_block.get()));
} else {
return Status::InternalError("maybe can't reach here, execute const
expr: {}, {}, {}",
_cur_child_id,
_get_first_materialized_child_idx(),
children_count());
}
- if (UNLIKELY(eos)) {
- //if _cur_child_id eos, need check to push block
- //Now here can't check _output_block rows, even it's row==0, also need
push block
- //because maybe sink is eos and queue have none data, if not push block
- //the source can't can_read again and can't set source finished
- if (local_state._output_block) {
- RETURN_IF_ERROR(local_state._shared_state->data_queue.push_block(
- std::move(local_state._output_block), _cur_child_id));
- }
- local_state._shared_state->data_queue.set_finish(_cur_child_id);
- return Status::OK();
- }
- // not eos and block rows is enough to output,so push block
- if (local_state._output_block && (local_state._output_block->rows() >=
state->batch_size())) {
- RETURN_IF_ERROR(local_state._shared_state->data_queue.push_block(
- std::move(local_state._output_block), _cur_child_id));
- }
- return Status::OK();
+ return
local_state._shared_state->data_queue.push_block(std::move(output_block),
_cur_child_id,
+ eos);
}
} // namespace doris
diff --git a/be/src/exec/operator/union_sink_operator.h
b/be/src/exec/operator/union_sink_operator.h
index 2bf0e7b3e75..7f2ec5bbb2f 100644
--- a/be/src/exec/operator/union_sink_operator.h
+++ b/be/src/exec/operator/union_sink_operator.h
@@ -51,8 +51,7 @@ class UnionSinkOperatorX;
class UnionSinkLocalState final : public
PipelineXSinkLocalState<UnionSharedState> {
public:
ENABLE_FACTORY_CREATOR(UnionSinkLocalState);
- UnionSinkLocalState(DataSinkOperatorXBase* parent, RuntimeState* state)
- : Base(parent, state), _child_row_idx(0) {}
+ UnionSinkLocalState(DataSinkOperatorXBase* parent, RuntimeState* state) :
Base(parent, state) {}
Status init(RuntimeState* state, LocalSinkStateInfo& info) override;
Status open(RuntimeState* state) override;
friend class UnionSinkOperatorX;
@@ -60,8 +59,6 @@ public:
using Parent = UnionSinkOperatorX;
private:
- std::unique_ptr<Block> _output_block;
-
/// Const exprs materialized by this node. These exprs don't refer to any
children.
/// Only materialized by the first fragment instance to avoid duplication.
VExprContextSPtrs _const_expr;
@@ -69,8 +66,6 @@ private:
/// Exprs materialized by this node. The i-th result expr list refers to
the i-th child.
VExprContextSPtrs _child_expr;
- /// Index of current row in child_row_block_.
- int _child_row_idx;
RuntimeProfile::Counter* _expr_timer = nullptr;
};
@@ -168,7 +163,6 @@ private:
RETURN_IF_ERROR(
materialize_block(local_state._child_expr,
input_block, &res, false));
}
- local_state._child_row_idx += res.rows();
RETURN_IF_ERROR(mblock.merge(res));
}
return Status::OK();
diff --git a/be/src/exec/operator/union_source_operator.cpp
b/be/src/exec/operator/union_source_operator.cpp
index 6cc39ebb7ce..926fa5487db 100644
--- a/be/src/exec/operator/union_source_operator.cpp
+++ b/be/src/exec/operator/union_source_operator.cpp
@@ -30,7 +30,6 @@
#include "exec/operator/union_sink_operator.h"
#include "exec/pipeline/dependency.h"
#include "runtime/descriptors.h"
-#include "util/defer_op.h"
namespace doris {
class RuntimeState;
@@ -93,49 +92,33 @@ std::string UnionSourceLocalState::debug_string(int
indentation_level) const {
fmt::memory_buffer debug_string_buffer;
fmt::format_to(debug_string_buffer, "{}",
Base::debug_string(indentation_level));
if (_shared_state) {
- fmt::format_to(debug_string_buffer, ", data_queue: (is_all_finish =
{}, has_data = {})",
- _shared_state->data_queue.is_all_finish(),
- _shared_state->data_queue.has_more_data());
+ fmt::format_to(debug_string_buffer, ", data_queue: {}",
+ _shared_state->data_queue.debug_string());
}
return fmt::to_string(debug_string_buffer);
}
Status UnionSourceOperatorX::get_block_impl(RuntimeState* state, Block* block,
bool* eos) {
auto& local_state = get_local_state(state);
- Defer set_eos {[&]() {
- // the eos check of union operator is complex, need check all logical
if you want modify
- // could ref this PR: https://github.com/apache/doris/pull/29677
- // have executing const expr, queue have no data anymore, and child
could be closed
- if (_child_size == 0 && !local_state._need_read_for_const_expr) {
- *eos = true;
- } else if (_has_data(state)) {
- *eos = false;
- } else if (local_state._shared_state->data_queue.is_all_finish()) {
- // Here, check the value of `_has_data(state)` again after
`data_queue.is_all_finish()` is TRUE
- // as there may be one or more blocks when
`data_queue.is_all_finish()` is TRUE.
- *eos = !_has_data(state);
- } else {
- *eos = false;
- }
- }};
-
+ *eos = false;
SCOPED_TIMER(local_state.exec_time_counter());
if (local_state._need_read_for_const_expr) {
if (has_more_const(state)) {
RETURN_IF_ERROR(get_next_const(state, block));
}
local_state._need_read_for_const_expr = has_more_const(state);
+ if (_child_size == 0 && !local_state._need_read_for_const_expr) {
+ *eos = true;
+ }
} else if (_child_size != 0) {
- std::unique_ptr<Block> output_block;
- int child_idx = 0;
-
RETURN_IF_ERROR(local_state._shared_state->data_queue.get_block_from_queue(&output_block,
-
&child_idx));
- if (!output_block) {
+ auto queue_block =
DORIS_TRY(local_state._shared_state->data_queue.get_block_from_queue());
+ *eos = queue_block.eos;
+ if (!queue_block.block) {
return Status::OK();
}
- block->swap(*output_block);
-
output_block->clear_column_data(row_descriptor().num_materialized_slots());
-
local_state._shared_state->data_queue.push_free_block(std::move(output_block),
child_idx);
+ block->swap(*queue_block.block);
+
queue_block.block->clear_column_data(row_descriptor().num_materialized_slots());
+
local_state._shared_state->data_queue.push_free_block(std::move(queue_block));
}
local_state.reached_limit(block, eos);
return Status::OK();
diff --git a/be/src/exec/operator/union_source_operator.h
b/be/src/exec/operator/union_source_operator.h
index c1006e46693..67b405b1002 100644
--- a/be/src/exec/operator/union_source_operator.h
+++ b/be/src/exec/operator/union_source_operator.h
@@ -122,13 +122,6 @@ public:
}
private:
- bool _has_data(RuntimeState* state) const {
- auto& local_state = get_local_state(state);
- if (_child_size == 0) {
- return local_state._need_read_for_const_expr;
- }
- return local_state._shared_state->data_queue.remaining_has_data();
- }
bool has_more_const(RuntimeState* state) const {
auto& local_state = get_local_state(state);
return state->per_fragment_instance_idx() == 0 &&
@@ -140,4 +133,4 @@ private:
std::vector<VExprContextSPtrs> _const_expr_lists;
};
-} // namespace doris
\ No newline at end of file
+} // namespace doris
diff --git a/be/test/exec/operator/union_operator_test.cpp
b/be/test/exec/operator/union_operator_test.cpp
index 3009e977a47..8590c86f732 100644
--- a/be/test/exec/operator/union_operator_test.cpp
+++ b/be/test/exec/operator/union_operator_test.cpp
@@ -254,7 +254,9 @@ TEST_F(UnionOperatorTest, test_sink_and_source) {
{
for (int i = 0; i < child_size; i++) {
- sink_state[i]->_batch_size = 2;
+ // Each sink input should be queued immediately, including
materialized children whose
+ // input is smaller than the runtime batch size.
+ sink_state[i]->_batch_size = 10;
Block block = ColumnHelper::create_block<DataTypeInt64>({1, 2},
{3, 4});
EXPECT_TRUE(sink_ops[i]->sink(sink_state[i].get(), &block, false));
}
@@ -291,4 +293,4 @@ TEST_F(UnionOperatorTest, test_sink_and_source) {
block, ColumnHelper::create_block<DataTypeInt64>({1, 2}, {3,
4})));
}
}
-} // namespace doris
\ No newline at end of file
+} // namespace doris
diff --git a/be/test/exec/pipeline/data_queue_test.cpp
b/be/test/exec/pipeline/data_queue_test.cpp
index f8ba26bc664..db1c3cccef4 100644
--- a/be/test/exec/pipeline/data_queue_test.cpp
+++ b/be/test/exec/pipeline/data_queue_test.cpp
@@ -221,57 +221,59 @@ public:
// Initial state: no data, no finish.
TEST_F(DataQueueTest, InitialState) {
- EXPECT_FALSE(data_queue->has_more_data());
- EXPECT_FALSE(data_queue->is_all_finish());
- EXPECT_FALSE(data_queue->remaining_has_data());
+ EXPECT_EQ(data_queue->debug_string(), "(is_all_finish = false, has_data =
false)");
+ auto queue_block = TEST_TRY(data_queue->get_block_from_queue());
+ EXPECT_EQ(queue_block.block, nullptr);
+ EXPECT_FALSE(queue_block.eos);
}
// Push one block and retrieve it.
TEST_F(DataQueueTest, SinglePushPop) {
- EXPECT_TRUE(data_queue->push_block(make_block(), 0).ok());
- EXPECT_TRUE(data_queue->has_more_data());
+ EXPECT_TRUE(data_queue->push_block(make_block(), 0, false).ok());
+ EXPECT_EQ(data_queue->debug_string(), "(is_all_finish = false, has_data =
true)");
- // Find the queue with data.
- EXPECT_TRUE(data_queue->remaining_has_data());
-
- std::unique_ptr<Block> out;
- int child_idx = -1;
- EXPECT_TRUE(data_queue->get_block_from_queue(&out, &child_idx).ok());
- EXPECT_NE(out, nullptr);
- EXPECT_EQ(child_idx, 0);
- EXPECT_FALSE(data_queue->has_more_data());
+ auto queue_block = TEST_TRY(data_queue->get_block_from_queue());
+ EXPECT_NE(queue_block.block, nullptr);
+ EXPECT_FALSE(queue_block.eos);
+ EXPECT_EQ(data_queue->debug_string(), "(is_all_finish = false, has_data =
false)");
}
-// is_all_finish only becomes true after all children call set_finish.
+// is_all_finish only becomes true after all children push eos.
TEST_F(DataQueueTest, IsAllFinishAfterAllChildren) {
- data_queue->set_finish(0);
- EXPECT_FALSE(data_queue->is_all_finish());
- data_queue->set_finish(1);
- EXPECT_FALSE(data_queue->is_all_finish());
- data_queue->set_finish(2);
- EXPECT_TRUE(data_queue->is_all_finish());
+ EXPECT_TRUE(data_queue->push_block(nullptr, 0, true).ok());
+ auto first_queue_block = TEST_TRY(data_queue->get_block_from_queue());
+ EXPECT_FALSE(first_queue_block.eos);
+ EXPECT_TRUE(data_queue->push_block(nullptr, 1, true).ok());
+ auto second_queue_block = TEST_TRY(data_queue->get_block_from_queue());
+ EXPECT_FALSE(second_queue_block.eos);
+ EXPECT_TRUE(data_queue->push_block(nullptr, 2, true).ok());
+ auto last_queue_block = TEST_TRY(data_queue->get_block_from_queue());
+ EXPECT_TRUE(last_queue_block.eos);
}
-// set_finish is idempotent.
-TEST_F(DataQueueTest, SetFinishIdempotent) {
- data_queue->set_finish(0);
- data_queue->set_finish(0); // second call must not double-decrement
- data_queue->set_finish(1);
- data_queue->set_finish(2);
- EXPECT_TRUE(data_queue->is_all_finish());
+// eos push is idempotent.
+TEST_F(DataQueueTest, EosPushIdempotent) {
+ EXPECT_TRUE(data_queue->push_block(nullptr, 0, true).ok());
+ EXPECT_TRUE(data_queue->push_block(nullptr, 0, true).ok());
+ EXPECT_TRUE(data_queue->push_block(nullptr, 1, true).ok());
+ EXPECT_TRUE(data_queue->push_block(nullptr, 2, true).ok());
+ auto queue_block = TEST_TRY(data_queue->get_block_from_queue());
+ EXPECT_TRUE(queue_block.eos);
}
-// child_idx returned by get_block_from_queue reflects the actual queue.
-TEST_F(DataQueueTest, ChildIdxReturned) {
+// DataQueueBlock can return a popped block to the originating child free list.
+TEST_F(DataQueueTest, ReturnBlockToChildFreeList) {
// Push to child 1 only.
- EXPECT_TRUE(data_queue->push_block(make_block(), 1).ok());
- data_queue->remaining_has_data(); // advance _flag_queue_idx to find child
1
+ EXPECT_TRUE(data_queue->push_block(make_block(), 1, false).ok());
- std::unique_ptr<Block> out;
- int child_idx = -1;
- EXPECT_TRUE(data_queue->get_block_from_queue(&out, &child_idx).ok());
- EXPECT_NE(out, nullptr);
- EXPECT_EQ(child_idx, 1);
+ auto queue_block = TEST_TRY(data_queue->get_block_from_queue());
+ EXPECT_NE(queue_block.block, nullptr);
+ queue_block.block->clear();
+ auto* returned_block = queue_block.block.get();
+ data_queue->push_free_block(std::move(queue_block));
+
+ auto reused_block = data_queue->get_free_block(1);
+ EXPECT_EQ(reused_block.get(), returned_block);
}
// get_free_block returns a new block when free list is empty, reuses when not.
@@ -282,7 +284,9 @@ TEST_F(DataQueueTest, FreeBlockReuse) {
// Return it to the free list.
block->clear(); // ensure rows == 0
- data_queue->push_free_block(std::move(block), 0);
+ DataQueueBlock queue_block;
+ queue_block.block = std::move(block);
+ data_queue->push_free_block(std::move(queue_block));
// Second call: must return the recycled block.
auto block2 = data_queue->get_free_block(0);
@@ -292,7 +296,9 @@ TEST_F(DataQueueTest, FreeBlockReuse) {
// In low-memory mode push_free_block discards blocks and max drops to 1.
TEST_F(DataQueueTest, LowMemoryMode) {
// Pre-populate the free list.
- data_queue->push_free_block(Block::create_unique(), 0);
+ DataQueueBlock queue_block;
+ queue_block.block = Block::create_unique();
+ data_queue->push_free_block(std::move(queue_block));
data_queue->set_low_memory_mode();
@@ -303,7 +309,9 @@ TEST_F(DataQueueTest, LowMemoryMode) {
// push_free_block now discards.
block->clear();
- data_queue->push_free_block(std::move(block), 0);
+ DataQueueBlock low_memory_block;
+ low_memory_block.block = std::move(block);
+ data_queue->push_free_block(std::move(low_memory_block));
auto block2 = data_queue->get_free_block(0);
// Still gets a fresh allocation (free list stays empty).
EXPECT_NE(block2, nullptr);
@@ -311,15 +319,17 @@ TEST_F(DataQueueTest, LowMemoryMode) {
// terminate() finishes all children and clears pending blocks from sub-queues.
TEST_F(DataQueueTest, Terminate) {
- EXPECT_TRUE(data_queue->push_block(make_block(), 0).ok());
- EXPECT_TRUE(data_queue->push_block(make_block(), 1).ok());
+ EXPECT_TRUE(data_queue->push_block(make_block(), 0, false).ok());
+ EXPECT_TRUE(data_queue->push_block(make_block(), 1, false).ok());
data_queue->terminate();
- EXPECT_TRUE(data_queue->is_all_finish());
- // remaining_has_data() checks blocks_in_queue per sub-queue,
- // which clear_blocks() resets to 0.
- EXPECT_FALSE(data_queue->remaining_has_data());
+ EXPECT_EQ(data_queue->debug_string(), "(is_all_finish = true, has_data =
false)");
+ source_dep->block();
+ EXPECT_TRUE(source_dep->ready());
+ auto queue_block = TEST_TRY(data_queue->get_block_from_queue());
+ EXPECT_EQ(queue_block.block, nullptr);
+ EXPECT_TRUE(queue_block.eos);
}
// set_max_blocks_in_sub_queue propagates to every sub-queue.
@@ -327,12 +337,12 @@ TEST_F(DataQueueTest, SetMaxBlocksInSubQueue) {
data_queue->set_max_blocks_in_sub_queue(5);
// Push 5 blocks to child 0 — sink must stay ready (not over the limit
yet).
for (int i = 0; i < 5; i++) {
- EXPECT_TRUE(data_queue->push_block(make_block(), 0).ok());
+ EXPECT_TRUE(data_queue->push_block(make_block(), 0, false).ok());
}
EXPECT_TRUE(sink_deps[0]->ready());
// 6th push exceeds limit → sink blocked.
- EXPECT_TRUE(data_queue->push_block(make_block(), 0).ok());
+ EXPECT_TRUE(data_queue->push_block(make_block(), 0, false).ok());
EXPECT_FALSE(sink_deps[0]->ready());
}
@@ -341,10 +351,28 @@ TEST_F(DataQueueTest, SourceReadyOnPush) {
source_dep->block(); // start blocked
EXPECT_FALSE(source_dep->ready());
- EXPECT_TRUE(data_queue->push_block(make_block(), 0).ok());
+ EXPECT_TRUE(data_queue->push_block(make_block(), 0, false).ok());
EXPECT_TRUE(source_dep->ready());
}
+TEST_F(DataQueueTest, SourceReadyOnEosOnly) {
+ source_dep->block();
+ EXPECT_FALSE(source_dep->ready());
+
+ EXPECT_TRUE(data_queue->push_block(nullptr, 0, true).ok());
+ EXPECT_TRUE(source_dep->ready());
+}
+
+TEST_F(DataQueueTest, EosOnlyAfterAllChildren) {
+ EXPECT_TRUE(data_queue->push_block(nullptr, 0, true).ok());
+ EXPECT_TRUE(data_queue->push_block(nullptr, 1, true).ok());
+ EXPECT_TRUE(data_queue->push_block(nullptr, 2, true).ok());
+
+ auto queue_block = TEST_TRY(data_queue->get_block_from_queue());
+ EXPECT_EQ(queue_block.block, nullptr);
+ EXPECT_TRUE(queue_block.eos);
+}
+
// ---------------------------------------------------------------------------
// Multi-threaded integration test (existing)
// ---------------------------------------------------------------------------
@@ -355,19 +383,9 @@ TEST_F(DataQueueTest, MultiTest) {
while (true) {
bool eos = false;
if (source_dep->ready()) {
- Defer set_eos {[&]() {
- if (data_queue->remaining_has_data()) {
- eos = false;
- } else if (data_queue->is_all_finish()) {
- eos = !data_queue->remaining_has_data();
- } else {
- eos = false;
- }
- }};
- std::unique_ptr<Block> output_block;
- int child_idx = 0;
- EXPECT_TRUE(data_queue->get_block_from_queue(&output_block,
&child_idx));
- if (output_block) {
+ auto queue_block =
TEST_TRY(data_queue->get_block_from_queue());
+ eos = queue_block.eos;
+ if (queue_block.block) {
output_count++;
}
}
@@ -392,11 +410,11 @@ TEST_F(DataQueueTest, MultiTest) {
int i = 0;
while (i < 50) {
if (sink_deps[id]->ready()) {
-
EXPECT_TRUE(data_queue->push_block(std::move(input_blocks[id][i]), id).ok());
+
EXPECT_TRUE(data_queue->push_block(std::move(input_blocks[id][i]), id,
false).ok());
i++;
}
}
- data_queue->set_finish(id);
+ EXPECT_TRUE(data_queue->push_block(nullptr, id, true).ok());
};
std::thread input1(input_func, 0);
@@ -409,7 +427,7 @@ TEST_F(DataQueueTest, MultiTest) {
output1.join();
EXPECT_EQ(output_count, 150);
- EXPECT_TRUE(data_queue->is_all_finish());
+ EXPECT_EQ(data_queue->debug_string(), "(is_all_finish = true, has_data =
false)");
}
// ./run-be-ut.sh --run --filter=DataQueueTest.*
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
