diff --git a/be/src/exec/pipeline/pipeline.cpp b/be/src/exec/pipeline/pipeline.cpp index bf2d3700df5b28..e79b133c5d8e74 100644 --- a/be/src/exec/pipeline/pipeline.cpp +++ b/be/src/exec/pipeline/pipeline.cpp @@ -19,8 +19,10 @@ #include #include +#include #include +#include "exec/exchange/local_exchange_sink_operator.h" #include "exec/exchange/local_exchange_source_operator.h" #include "exec/operator/operator.h" #include "exec/pipeline/pipeline_fragment_context.h" @@ -28,6 +30,51 @@ namespace doris { +Status validate_paired_pipeline_task_count(const Pipelines& pipelines) { + struct OperatorLocation { + const Pipeline* pipeline; + const OperatorXBase* op; + }; + + std::unordered_map operator_locations; + for (const auto& pipeline : pipelines) { + for (const auto& op : pipeline->operators()) { + const bool inserted = + operator_locations + .emplace(op->operator_id(), OperatorLocation {pipeline.get(), op.get()}) + .second; + DORIS_CHECK(inserted) << "duplicate operator id " << op->operator_id(); + } + } + + for (const auto& pipeline : pipelines) { + const auto* sink = pipeline->sink(); + DORIS_CHECK(sink != nullptr) << pipeline->debug_string(); + if (dynamic_cast(sink) != nullptr) { + continue; + } + for (const auto dest_id : sink->dests_id()) { + const auto destination = operator_locations.find(dest_id); + // Final data sinks target operators outside this fragment and have no local pair. + if (destination == operator_locations.end()) { + continue; + } + const auto* destination_pipeline = destination->second.pipeline; + if (pipeline->num_tasks() != destination_pipeline->num_tasks()) { + return Status::InternalError( + "Non-local-exchange pipeline task count mismatch: sink {} (operator id {}, " + "pipeline id {}, task count {}) targets {} (operator id {}, pipeline id " + "{}, " + "task count {})", + sink->get_name(), sink->operator_id(), pipeline->id(), + pipeline->num_tasks(), destination->second.op->get_name(), dest_id, + destination_pipeline->id(), destination_pipeline->num_tasks()); + } + } + } + return Status::OK(); +} + void Pipeline::_init_profile() { auto s = fmt::format("Pipeline (pipeline id={})", _pipeline_id); _pipeline_profile = std::make_unique(std::move(s)); diff --git a/be/src/exec/pipeline/pipeline.h b/be/src/exec/pipeline/pipeline.h index 17fdecbd52c66b..2f1198024ca54b 100644 --- a/be/src/exec/pipeline/pipeline.h +++ b/be/src/exec/pipeline/pipeline.h @@ -177,4 +177,8 @@ class Pipeline : public std::enable_shared_from_this { // Parallelism of parent pipeline. const int _num_tasks_of_parent; }; + +// Pipeline-breaking sink/destination pairs share per-task state and require one-to-one task +// counts. Local exchange is the only pair allowed to redistribute between different counts. +Status validate_paired_pipeline_task_count(const Pipelines& pipelines); } // namespace doris diff --git a/be/src/exec/pipeline/pipeline_fragment_context.cpp b/be/src/exec/pipeline/pipeline_fragment_context.cpp index f62bd0730816c2..34886ffba1ae10 100644 --- a/be/src/exec/pipeline/pipeline_fragment_context.cpp +++ b/be/src/exec/pipeline/pipeline_fragment_context.cpp @@ -338,6 +338,7 @@ Status PipelineFragmentContext::_build_and_prepare_full_pipeline(ThreadPool* thr _params.bucket_seq_to_instance_idx, _params.shuffle_idx_to_instance_idx)); } + RETURN_IF_ERROR(validate_paired_pipeline_task_count(_pipelines)); // 5. Initialize global states in pipelines. for (PipelinePtr& pipeline : _pipelines) { diff --git a/be/test/exec/pipeline/pipeline_test.cpp b/be/test/exec/pipeline/pipeline_test.cpp index 4cd255648b6219..202e1e54137bca 100644 --- a/be/test/exec/pipeline/pipeline_test.cpp +++ b/be/test/exec/pipeline/pipeline_test.cpp @@ -27,6 +27,7 @@ #include "core/column/column_vector.h" #include "core/data_type/data_type.h" #include "core/data_type/data_type_number.h" +#include "exec/exchange/local_exchange_sink_operator.h" #include "exec/exchange/vdata_stream_mgr.h" #include "exec/operator/exchange_source_operator.h" #include "exec/operator/hashjoin_build_sink.h" @@ -48,6 +49,49 @@ namespace doris { static void empty_function(RuntimeState*, Status*) {} +class PipelineTaskCountTestOperator final : public OperatorXBase { +public: + explicit PipelineTaskCountTestOperator(int operator_id) + : OperatorXBase(nullptr, operator_id, operator_id) {} + + Status get_block_impl(RuntimeState* state, Block* block, bool* eos) override { + return Status::OK(); + } + + Status setup_local_state(RuntimeState* state, LocalStateInfo& info) override { + return Status::OK(); + } +}; + +class PipelineTaskCountTestSinkOperator final : public DataSinkOperatorXBase { +public: + PipelineTaskCountTestSinkOperator(int operator_id, std::vector dest_ids) + : DataSinkOperatorXBase(operator_id, operator_id, dest_ids) {} + + Status sink_impl(RuntimeState* state, Block* block, bool eos) override { return Status::OK(); } + + Status setup_local_state(RuntimeState* state, LocalSinkStateInfo& info) override { + return Status::OK(); + } + + std::shared_ptr create_shared_state() const override { return nullptr; } +}; + +static PipelinePtr create_pipeline_for_task_count_test(int pipeline_id, int num_tasks, + int operator_id) { + auto pipeline = std::make_shared(pipeline_id, num_tasks, num_tasks); + OperatorPtr op = std::make_shared(operator_id); + EXPECT_TRUE(pipeline->add_operator(op, 0).ok()); + return pipeline; +} + +static void set_sink_for_task_count_test(const PipelinePtr& pipeline, int sink_id, + std::vector dest_ids) { + DataSinkOperatorPtr sink = + std::make_shared(sink_id, std::move(dest_ids)); + EXPECT_TRUE(pipeline->set_sink(sink).ok()); +} + class PipelineTest : public testing::Test { public: PipelineTest() : _obj_pool(new ObjectPool()), _mgr(std::make_unique()) {} @@ -1292,4 +1336,46 @@ TEST_F(PipelineTest, QueryTaskProgressCountersSurviveReset) { EXPECT_EQ(ctrl2->get_finished_task_num(), 0); } +TEST(PipelineTaskCountTest, ValidatePairedPipelineTaskCount) { + auto destination = create_pipeline_for_task_count_test(0, 3, 10); + set_sink_for_task_count_test(destination, -100, {1000}); + auto matching_sink = create_pipeline_for_task_count_test(1, 3, 11); + set_sink_for_task_count_test(matching_sink, -1, {10}); + + EXPECT_TRUE(validate_paired_pipeline_task_count({destination, matching_sink}).ok()); + + auto mismatched_sink = create_pipeline_for_task_count_test(2, 1, 12); + set_sink_for_task_count_test(mismatched_sink, -2, {10}); + auto status = validate_paired_pipeline_task_count({destination, mismatched_sink}); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("pipeline id 2, task count 1"), std::string::npos); + EXPECT_NE(status.to_string().find("pipeline id 0, task count 3"), std::string::npos); +} + +TEST(PipelineTaskCountTest, AllowLocalExchangeTaskCountMismatch) { + auto destination = create_pipeline_for_task_count_test(0, 3, 10); + set_sink_for_task_count_test(destination, -100, {1000}); + auto local_exchange_sink = create_pipeline_for_task_count_test(1, 1, 11); + DataSinkOperatorPtr sink = std::make_shared( + -1, 10, 3, std::vector {}, std::map {}); + EXPECT_TRUE(local_exchange_sink->set_sink(sink).ok()); + + EXPECT_TRUE(validate_paired_pipeline_task_count({destination, local_exchange_sink}).ok()); +} + +TEST(PipelineTaskCountTest, ValidateEverySinkDestination) { + auto first_destination = create_pipeline_for_task_count_test(0, 3, 10); + set_sink_for_task_count_test(first_destination, -100, {1000}); + auto second_destination = create_pipeline_for_task_count_test(1, 1, 20); + set_sink_for_task_count_test(second_destination, -200, {2000}); + auto multi_destination_sink = create_pipeline_for_task_count_test(2, 3, 30); + set_sink_for_task_count_test(multi_destination_sink, -1, {10, 20}); + + auto status = validate_paired_pipeline_task_count( + {first_destination, second_destination, multi_destination_sink}); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("operator id 20"), std::string::npos); + EXPECT_NE(status.to_string().find("pipeline id 1, task count 1"), std::string::npos); +} + } // namespace doris