// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.

#include "exec/pipeline/pipeline_fragment_context.h"

#include <gen_cpp/DataSinks_types.h>

#include <gen_cpp/FrontendService.h>

#include <gen_cpp/FrontendService_types.h>

#include <gen_cpp/PaloInternalService_types.h>

#include <gen_cpp/PlanNodes_types.h>

#include <pthread.h>

#include <algorithm>

#include <cstdlib>

// IWYU pragma: no_include <bits/chrono.h>

#include <fmt/format.h>

#include <thrift/Thrift.h>

#include <thrift/protocol/TDebugProtocol.h>

#include <thrift/transport/TTransportException.h>

#include <chrono> // IWYU pragma: keep

#include <map>

#include <memory>

#include <ostream>

#include <utility>

#include "cloud/config.h"

#include "common/cast_set.h"

#include "common/config.h"

#include "common/exception.h"

#include "common/logging.h"

#include "common/status.h"

#include "exec/exchange/local_exchange_sink_operator.h"

#include "exec/exchange/local_exchange_source_operator.h"

#include "exec/exchange/local_exchanger.h"

#include "exec/exchange/vdata_stream_mgr.h"

#include "exec/operator/aggregation_sink_operator.h"

#include "exec/operator/aggregation_source_operator.h"

#include "exec/operator/analytic_sink_operator.h"

#include "exec/operator/analytic_source_operator.h"

#include "exec/operator/assert_num_rows_operator.h"

#include "exec/operator/blackhole_sink_operator.h"

#include "exec/operator/bucketed_aggregation_sink_operator.h"

#include "exec/operator/bucketed_aggregation_source_operator.h"

#include "exec/operator/cache_sink_operator.h"

#include "exec/operator/cache_source_operator.h"

#include "exec/operator/datagen_operator.h"

#include "exec/operator/dict_sink_operator.h"

#include "exec/operator/distinct_streaming_aggregation_operator.h"

#include "exec/operator/empty_set_operator.h"

#include "exec/operator/exchange_sink_operator.h"

#include "exec/operator/exchange_source_operator.h"

#include "exec/operator/file_scan_operator.h"

#include "exec/operator/group_commit_block_sink_operator.h"

#include "exec/operator/group_commit_scan_operator.h"

#include "exec/operator/hashjoin_build_sink.h"

#include "exec/operator/hashjoin_probe_operator.h"

#include "exec/operator/hive_table_sink_operator.h"

#include "exec/operator/iceberg_delete_sink_operator.h"

#include "exec/operator/iceberg_merge_sink_operator.h"

#include "exec/operator/iceberg_table_sink_operator.h"

#include "exec/operator/jdbc_scan_operator.h"

#include "exec/operator/jdbc_table_sink_operator.h"

#include "exec/operator/local_merge_sort_source_operator.h"

#include "exec/operator/materialization_opertor.h"

#include "exec/operator/maxcompute_table_sink_operator.h"

#include "exec/operator/memory_scratch_sink_operator.h"

#include "exec/operator/meta_scan_operator.h"

#include "exec/operator/multi_cast_data_stream_sink.h"

#include "exec/operator/multi_cast_data_stream_source.h"

#include "exec/operator/nested_loop_join_build_operator.h"

#include "exec/operator/nested_loop_join_probe_operator.h"

#include "exec/operator/olap_scan_operator.h"

#include "exec/operator/olap_table_sink_operator.h"

#include "exec/operator/olap_table_sink_v2_operator.h"

#include "exec/operator/partition_sort_sink_operator.h"

#include "exec/operator/partition_sort_source_operator.h"

#include "exec/operator/partitioned_aggregation_sink_operator.h"

#include "exec/operator/partitioned_aggregation_source_operator.h"

#include "exec/operator/partitioned_hash_join_probe_operator.h"

#include "exec/operator/partitioned_hash_join_sink_operator.h"

#include "exec/operator/rec_cte_anchor_sink_operator.h"

#include "exec/operator/rec_cte_scan_operator.h"

#include "exec/operator/rec_cte_sink_operator.h"

#include "exec/operator/rec_cte_source_operator.h"

#include "exec/operator/repeat_operator.h"

#include "exec/operator/result_file_sink_operator.h"

#include "exec/operator/result_sink_operator.h"

#include "exec/operator/schema_scan_operator.h"

#include "exec/operator/select_operator.h"

#include "exec/operator/set_probe_sink_operator.h"

#include "exec/operator/set_sink_operator.h"

#include "exec/operator/set_source_operator.h"

#include "exec/operator/sort_sink_operator.h"

#include "exec/operator/sort_source_operator.h"

#include "exec/operator/spill_iceberg_table_sink_operator.h"

#include "exec/operator/spill_sort_sink_operator.h"

#include "exec/operator/spill_sort_source_operator.h"

#include "exec/operator/streaming_aggregation_operator.h"

#include "exec/operator/table_function_operator.h"

#include "exec/operator/tvf_table_sink_operator.h"

#include "exec/operator/union_sink_operator.h"

#include "exec/operator/union_source_operator.h"

#include "exec/pipeline/dependency.h"

#include "exec/pipeline/pipeline_task.h"

#include "exec/pipeline/task_scheduler.h"

#include "exec/runtime_filter/runtime_filter_mgr.h"

#include "exec/sort/topn_sorter.h"

#include "exec/spill/spill_file.h"

#include "io/fs/stream_load_pipe.h"

#include "load/stream_load/new_load_stream_mgr.h"

#include "runtime/exec_env.h"

#include "runtime/fragment_mgr.h"

#include "runtime/result_buffer_mgr.h"

#include "runtime/runtime_state.h"

#include "runtime/thread_context.h"

#include "service/backend_options.h"

#include "util/client_cache.h"

#include "util/countdown_latch.h"

#include "util/debug_util.h"

#include "util/network_util.h"

#include "util/uid_util.h"

namespace doris {

PipelineFragmentContext::PipelineFragmentContext(

        TUniqueId query_id, const TPipelineFragmentParams& request,

        std::shared_ptr<QueryContext> query_ctx, ExecEnv* exec_env,

        const std::function<void(RuntimeState*, Status*)>& call_back)

        : _query_id(std::move(query_id)),

          _fragment_id(request.fragment_id),

          _exec_env(exec_env),

          _query_ctx(std::move(query_ctx)),

          _call_back(call_back),

          _is_report_on_cancel(true),

          _params(request),

          _parallel_instances(_params.__isset.parallel_instances ? _params.parallel_instances : 0),

          _need_notify_close(request.__isset.need_notify_close ? request.need_notify_close

                                                               : false) {

    _fragment_watcher.start();

}

PipelineFragmentContext::~PipelineFragmentContext() {

    LOG_INFO("PipelineFragmentContext::~PipelineFragmentContext")

            .tag("query_id", print_id(_query_id))

            .tag("fragment_id", _fragment_id);

    _release_resource();

    {

        // The memory released by the query end is recorded in the query mem tracker.

        SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(_query_ctx->query_mem_tracker());

        _runtime_state.reset();

        _query_ctx.reset();

    }

}

bool PipelineFragmentContext::is_timeout(timespec now) const {

    if (_timeout <= 0) {

        return false;

    }

    return _fragment_watcher.elapsed_time_seconds(now) > _timeout;

}

// notify_close() transitions the PFC from "waiting for external close notification" to

// "self-managed close". For recursive CTE fragments, the old PFC is kept alive until

// the rerun_fragment(wait_for_destroy) RPC calls this to trigger shutdown.

// Returns true if all tasks have already closed (i.e., the PFC can be safely destroyed).

bool PipelineFragmentContext::notify_close() {

    bool all_closed = false;

    bool need_remove = false;

    {

        std::lock_guard<std::mutex> l(_task_mutex);

        if (_closed_tasks >= _total_tasks) {

            if (_need_notify_close) {

                // Fragment was cancelled and waiting for notify to close.

                // Record that we need to remove from fragment mgr, but do it

                // after releasing _task_mutex to avoid ABBA deadlock with

                // dump_pipeline_tasks() (which acquires _pipeline_map lock

                // first, then _task_mutex via debug_string()).

                need_remove = true;

            }

            all_closed = true;

        }

        // make fragment release by self after cancel

        _need_notify_close = false;

    }

    if (need_remove) {

        _exec_env->fragment_mgr()->remove_pipeline_context({_query_id, _fragment_id});

    }

    return all_closed;

}

// Must not add lock in this method. Because it will call query ctx cancel. And

// QueryCtx cancel will call fragment ctx cancel. And Also Fragment ctx's running

// Method like exchange sink buffer will call query ctx cancel. If we add lock here

// There maybe dead lock.

void PipelineFragmentContext::cancel(const Status reason) {

    LOG_INFO("PipelineFragmentContext::cancel")

            .tag("query_id", print_id(_query_id))

            .tag("fragment_id", _fragment_id)

            .tag("reason", reason.to_string());

    if (notify_close()) {

        return;

    }

    // Timeout is a special error code, we need print current stack to debug timeout issue.

    if (reason.is<ErrorCode::TIMEOUT>()) {

        auto dbg_str = fmt::format("PipelineFragmentContext is cancelled due to timeout:\n{}",

                                   debug_string());

        LOG_LONG_STRING(WARNING, dbg_str);

    }

    // `ILLEGAL_STATE` means queries this fragment belongs to was not found in FE (maybe finished)

    if (reason.is<ErrorCode::ILLEGAL_STATE>()) {

        LOG_WARNING("PipelineFragmentContext is cancelled due to illegal state : {}",

                    debug_string());

    }

    if (reason.is<ErrorCode::MEM_LIMIT_EXCEEDED>() || reason.is<ErrorCode::MEM_ALLOC_FAILED>()) {

        print_profile("cancel pipeline, reason: " + reason.to_string());

    }

    if (auto error_url = get_load_error_url(); !error_url.empty()) {

        _query_ctx->set_load_error_url(error_url);

    }

    if (auto first_error_msg = get_first_error_msg(); !first_error_msg.empty()) {

        _query_ctx->set_first_error_msg(first_error_msg);

    }

    _query_ctx->cancel(reason, _fragment_id);

    if (reason.is<ErrorCode::LIMIT_REACH>()) {

        _is_report_on_cancel = false;

    } else {

        for (auto& id : _fragment_instance_ids) {

            LOG(WARNING) << "PipelineFragmentContext cancel instance: " << print_id(id);

        }

    }

    // Get pipe from new load stream manager and send cancel to it or the fragment may hang to wait read from pipe

    // For stream load the fragment's query_id == load id, it is set in FE.

    auto stream_load_ctx = _exec_env->new_load_stream_mgr()->get(_query_id);

    if (stream_load_ctx != nullptr) {

        stream_load_ctx->pipe->cancel(reason.to_string());

        // Set error URL here because after pipe is cancelled, stream load execution may return early.

        // We need to set the error URL at this point to ensure error information is properly

        // propagated to the client.

        stream_load_ctx->error_url = get_load_error_url();

        stream_load_ctx->first_error_msg = get_first_error_msg();

    }

    for (auto& tasks : _tasks) {

        for (auto& task : tasks) {

            task.first->unblock_all_dependencies();

        }

    }

}

PipelinePtr PipelineFragmentContext::add_pipeline(PipelinePtr parent, int idx) {

    PipelineId id = _next_pipeline_id++;

    auto pipeline = std::make_shared<Pipeline>(

            id, parent ? std::min(parent->num_tasks(), _num_instances) : _num_instances,

            parent ? parent->num_tasks() : _num_instances);

    if (idx >= 0) {

        _pipelines.insert(_pipelines.begin() + idx, pipeline);

    } else {

        _pipelines.emplace_back(pipeline);

    }

    if (parent) {

        parent->set_children(pipeline);

    }

    return pipeline;

}

Status PipelineFragmentContext::_build_and_prepare_full_pipeline(ThreadPool* thread_pool) {

    {

        SCOPED_TIMER(_build_pipelines_timer);

        // 2. Build pipelines with operators in this fragment.

        auto root_pipeline = add_pipeline();

        RETURN_IF_ERROR(_build_pipelines(_runtime_state->obj_pool(), *_query_ctx->desc_tbl,

                                         &_root_op, root_pipeline));

        // Propagate _num_instances from LOCAL_EXCHANGE pipelines to ancestor pipelines

        // that inherited reduced num_tasks from a serial operator.

        _propagate_local_exchange_num_tasks();

        // Create deferred local exchangers now that all pipelines have final num_tasks.

        RETURN_IF_ERROR(_create_deferred_local_exchangers());

        // Raise num_tasks for pipelines whose serial non-scan operators (e.g.,

        // UNPARTITIONED Exchange) reduced num_tasks below _num_instances.

        // Without this, fragment instances 1+ have no task for these pipelines

        // and downstream operators fail with "must set shared state".

        //

        // This applies to ALL pipelines (not just deferred exchanger upstreams):

        // fragments with UNION/INTERSECT/EXCEPT + serial Exchange in child

        // pipelines also need the raise, even without FE-planned local exchange.

        //

        // Exception: serial scan sources (pooling scan) keep num_tasks=1 — the

        // PassthroughExchanger(1, N) handles the fan-out correctly.

        // NOTE: Do NOT raise pipelines whose source is a serial operator

        // (Exchange or scan) — they legitimately have 1 task, and raising

        // them causes crashes (e.g., 4 Exchange tasks but only 1 receives

        // data).  The correct fix for shared state injection across

        // instances is handled by the FE: it inserts local exchange nodes

        // between serial operators and their downstream consumers, creating

        // proper pipeline boundaries with _num_instances tasks.

        // 3. Create sink operator

        if (!_params.fragment.__isset.output_sink) {

            return Status::InternalError("No output sink in this fragment!");

        }

        RETURN_IF_ERROR(_create_data_sink(_runtime_state->obj_pool(), _params.fragment.output_sink,

                                          _params.fragment.output_exprs, _params,

                                          root_pipeline->output_row_desc(), _runtime_state.get(),

                                          *_desc_tbl, root_pipeline->id()));

        RETURN_IF_ERROR(_sink->init(_params.fragment.output_sink));

        RETURN_IF_ERROR(root_pipeline->set_sink(_sink));

        for (PipelinePtr& pipeline : _pipelines) {

            DCHECK(pipeline->sink() != nullptr) << pipeline->operators().size();

            RETURN_IF_ERROR(pipeline->sink()->set_child(pipeline->operators().back()));

        }

    }

    // 4. Build local exchanger

    if (_runtime_state->plan_local_shuffle()) {

        SCOPED_TIMER(_plan_local_exchanger_timer);

        RETURN_IF_ERROR(_plan_local_exchange(_params.num_buckets,

                                             _params.bucket_seq_to_instance_idx,

                                             _params.shuffle_idx_to_instance_idx));

    }

    // 5. Initialize global states in pipelines.

    for (PipelinePtr& pipeline : _pipelines) {

        SCOPED_TIMER(_prepare_all_pipelines_timer);

        pipeline->children().clear();

        RETURN_IF_ERROR(pipeline->prepare(_runtime_state.get()));

    }

    {

        SCOPED_TIMER(_build_tasks_timer);

        // 6. Build pipeline tasks and initialize local state.

        RETURN_IF_ERROR(_build_pipeline_tasks(thread_pool));

    }

    return Status::OK();

}

Status PipelineFragmentContext::prepare(ThreadPool* thread_pool) {

    if (_prepared) {

        return Status::InternalError("Already prepared");

    }

    if (_params.__isset.query_options && _params.query_options.__isset.execution_timeout) {

        _timeout = _params.query_options.execution_timeout;

    }

    _fragment_level_profile = std::make_unique<RuntimeProfile>("PipelineContext");

    _prepare_timer = ADD_TIMER(_fragment_level_profile, "PrepareTime");

    SCOPED_TIMER(_prepare_timer);

    _build_pipelines_timer = ADD_TIMER(_fragment_level_profile, "BuildPipelinesTime");

    _init_context_timer = ADD_TIMER(_fragment_level_profile, "InitContextTime");

    _plan_local_exchanger_timer = ADD_TIMER(_fragment_level_profile, "PlanLocalLocalExchangerTime");

    _build_tasks_timer = ADD_TIMER(_fragment_level_profile, "BuildTasksTime");

    _prepare_all_pipelines_timer = ADD_TIMER(_fragment_level_profile, "PrepareAllPipelinesTime");

    {

        SCOPED_TIMER(_init_context_timer);

        cast_set(_num_instances, _params.local_params.size());

        _total_instances =

                _params.__isset.total_instances ? _params.total_instances : _num_instances;

        auto* fragment_context = this;

        if (_params.query_options.__isset.is_report_success) {

            fragment_context->set_is_report_success(_params.query_options.is_report_success);

        }

        // 1. Set up the global runtime state.

        _runtime_state = RuntimeState::create_unique(

                _params.query_id, _params.fragment_id, _params.query_options,

                _query_ctx->query_globals, _exec_env, _query_ctx.get());

        _runtime_state->set_task_execution_context(shared_from_this());

        SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(_runtime_state->query_mem_tracker());

        if (_params.__isset.backend_id) {

            _runtime_state->set_backend_id(_params.backend_id);

        }

        if (_params.__isset.import_label) {

            _runtime_state->set_import_label(_params.import_label);

        }

        if (_params.__isset.db_name) {

            _runtime_state->set_db_name(_params.db_name);

        }

        if (_params.__isset.load_job_id) {

            _runtime_state->set_load_job_id(_params.load_job_id);

        }

        if (_params.is_simplified_param) {

            _desc_tbl = _query_ctx->desc_tbl;

        } else {

            DCHECK(_params.__isset.desc_tbl);

            RETURN_IF_ERROR(DescriptorTbl::create(_runtime_state->obj_pool(), _params.desc_tbl,

                                                  &_desc_tbl));

        }

        _runtime_state->set_desc_tbl(_desc_tbl);

        _runtime_state->set_num_per_fragment_instances(_params.num_senders);

        _runtime_state->set_load_stream_per_node(_params.load_stream_per_node);

        _runtime_state->set_total_load_streams(_params.total_load_streams);

        _runtime_state->set_num_local_sink(_params.num_local_sink);

        // init fragment_instance_ids

        const auto target_size = _params.local_params.size();

        _fragment_instance_ids.resize(target_size);

        for (size_t i = 0; i < _params.local_params.size(); i++) {

            auto fragment_instance_id = _params.local_params[i].fragment_instance_id;

            _fragment_instance_ids[i] = fragment_instance_id;

        }

    }

    RETURN_IF_ERROR(_build_and_prepare_full_pipeline(thread_pool));

    _init_next_report_time();

    _prepared = true;

    return Status::OK();

}

Status PipelineFragmentContext::_build_pipeline_tasks_for_instance(

        int instance_idx,

        const std::vector<std::shared_ptr<RuntimeProfile>>& pipeline_id_to_profile) {

    const auto& local_params = _params.local_params[instance_idx];

    auto fragment_instance_id = local_params.fragment_instance_id;

    auto runtime_filter_mgr = std::make_unique<RuntimeFilterMgr>(false);

    std::map<PipelineId, PipelineTask*> pipeline_id_to_task;

    auto get_shared_state = [&](PipelinePtr pipeline)

            -> std::map<int, std::pair<std::shared_ptr<BasicSharedState>,

                                       std::vector<std::shared_ptr<Dependency>>>> {

        std::map<int, std::pair<std::shared_ptr<BasicSharedState>,

                                std::vector<std::shared_ptr<Dependency>>>>

                shared_state_map;

        for (auto& op : pipeline->operators()) {

            auto source_id = op->operator_id();

            if (auto iter = _op_id_to_shared_state.find(source_id);

                iter != _op_id_to_shared_state.end()) {

                shared_state_map.insert({source_id, iter->second});

            }

        }

        for (auto sink_to_source_id : pipeline->sink()->dests_id()) {

            if (auto iter = _op_id_to_shared_state.find(sink_to_source_id);

                iter != _op_id_to_shared_state.end()) {

                shared_state_map.insert({sink_to_source_id, iter->second});

            }

        }

        return shared_state_map;

    };

    for (size_t pip_idx = 0; pip_idx < _pipelines.size(); pip_idx++) {

        auto& pipeline = _pipelines[pip_idx];

        if (pipeline->num_tasks() > 1 || instance_idx == 0) {

            auto task_runtime_state = RuntimeState::create_unique(

                    local_params.fragment_instance_id, _params.query_id, _params.fragment_id,

                    _params.query_options, _query_ctx->query_globals, _exec_env, _query_ctx.get());

            {

                // Initialize runtime state for this task

                task_runtime_state->set_query_mem_tracker(_query_ctx->query_mem_tracker());

                task_runtime_state->set_task_execution_context(shared_from_this());

                task_runtime_state->set_be_number(local_params.backend_num);

                if (_params.__isset.backend_id) {

                    task_runtime_state->set_backend_id(_params.backend_id);

                }

                if (_params.__isset.import_label) {

                    task_runtime_state->set_import_label(_params.import_label);

                }

                if (_params.__isset.db_name) {

                    task_runtime_state->set_db_name(_params.db_name);

                }

                if (_params.__isset.load_job_id) {

                    task_runtime_state->set_load_job_id(_params.load_job_id);

                }

                if (_params.__isset.wal_id) {

                    task_runtime_state->set_wal_id(_params.wal_id);

                }

                if (_params.__isset.content_length) {

                    task_runtime_state->set_content_length(_params.content_length);

                }

                task_runtime_state->set_desc_tbl(_desc_tbl);

                task_runtime_state->set_per_fragment_instance_idx(local_params.sender_id);

                task_runtime_state->set_num_per_fragment_instances(_params.num_senders);

                task_runtime_state->resize_op_id_to_local_state(max_operator_id());

                task_runtime_state->set_max_operator_id(max_operator_id());

                task_runtime_state->set_load_stream_per_node(_params.load_stream_per_node);

                task_runtime_state->set_total_load_streams(_params.total_load_streams);

                task_runtime_state->set_num_local_sink(_params.num_local_sink);

                task_runtime_state->set_runtime_filter_mgr(runtime_filter_mgr.get());

            }

            auto cur_task_id = _total_tasks++;

            task_runtime_state->set_task_id(cur_task_id);

            task_runtime_state->set_task_num(pipeline->num_tasks());

            auto task = std::make_shared<PipelineTask>(

                    pipeline, cur_task_id, task_runtime_state.get(),

                    std::dynamic_pointer_cast<PipelineFragmentContext>(shared_from_this()),

                    pipeline_id_to_profile[pip_idx].get(), get_shared_state(pipeline),

                    instance_idx);

            pipeline->incr_created_tasks(instance_idx, task.get());

            pipeline_id_to_task.insert({pipeline->id(), task.get()});

            _tasks[instance_idx].emplace_back(

                    std::pair<std::shared_ptr<PipelineTask>, std::unique_ptr<RuntimeState>> {

                            std::move(task), std::move(task_runtime_state)});

        }

    }

    /**

         * Build DAG for pipeline tasks.

         * For example, we have

         *

         *   ExchangeSink (Pipeline1)     JoinBuildSink (Pipeline2)

         *            \                      /

         *          JoinProbeOperator1 (Pipeline1)    JoinBuildSink (Pipeline3)

         *                 \                          /

         *               JoinProbeOperator2 (Pipeline1)

         *

         * In this fragment, we have three pipelines and pipeline 1 depends on pipeline 2 and pipeline 3.

         * To build this DAG, `_dag` manage dependencies between pipelines by pipeline ID and

         * `pipeline_id_to_task` is used to find the task by a unique pipeline ID.

         *

         * Finally, we have two upstream dependencies in Pipeline1 corresponding to JoinProbeOperator1

         * and JoinProbeOperator2.

         */

    for (auto& _pipeline : _pipelines) {

        if (pipeline_id_to_task.contains(_pipeline->id())) {

            auto* task = pipeline_id_to_task[_pipeline->id()];

            DCHECK(task != nullptr);

            // If this task has upstream dependency, then inject it into this task.

            if (_dag.contains(_pipeline->id())) {

                auto& deps = _dag[_pipeline->id()];

                for (auto& dep : deps) {

                    if (pipeline_id_to_task.contains(dep)) {

                        auto ss = pipeline_id_to_task[dep]->get_sink_shared_state();

                        if (ss) {

                            task->inject_shared_state(ss);

                        } else {

                            pipeline_id_to_task[dep]->inject_shared_state(

                                    task->get_source_shared_state());

                        }

                    }

                }

            }

        }

    }

    for (size_t pip_idx = 0; pip_idx < _pipelines.size(); pip_idx++) {

        if (pipeline_id_to_task.contains(_pipelines[pip_idx]->id())) {

            auto* task = pipeline_id_to_task[_pipelines[pip_idx]->id()];

            DCHECK(pipeline_id_to_profile[pip_idx]);

            std::vector<TScanRangeParams> scan_ranges;

            auto node_id = _pipelines[pip_idx]->operators().front()->node_id();

            if (local_params.per_node_scan_ranges.contains(node_id)) {

                scan_ranges = local_params.per_node_scan_ranges.find(node_id)->second;

            }

            RETURN_IF_ERROR_OR_CATCH_EXCEPTION(task->prepare(scan_ranges, local_params.sender_id,

                                                             _params.fragment.output_sink));

        }

    }

    {

        std::lock_guard<std::mutex> l(_state_map_lock);

        _runtime_filter_mgr_map[instance_idx] = std::move(runtime_filter_mgr);

    }

    return Status::OK();

}

Status PipelineFragmentContext::_build_pipeline_tasks(ThreadPool* thread_pool) {

    _total_tasks = 0;

    _closed_tasks = 0;

    const auto target_size = _params.local_params.size();

    _tasks.resize(target_size);

    _runtime_filter_mgr_map.resize(target_size);

    for (size_t pip_idx = 0; pip_idx < _pipelines.size(); pip_idx++) {

        _pip_id_to_pipeline[_pipelines[pip_idx]->id()] = _pipelines[pip_idx].get();

    }

    auto pipeline_id_to_profile = _runtime_state->build_pipeline_profile(_pipelines.size());

    if (target_size > 1 &&

        (_runtime_state->query_options().__isset.parallel_prepare_threshold &&

         target_size > _runtime_state->query_options().parallel_prepare_threshold)) {

        // If instances parallelism is big enough ( > parallel_prepare_threshold), we will prepare all tasks by multi-threads

        std::vector<Status> prepare_status(target_size);

        int submitted_tasks = 0;

        Status submit_status;

        CountDownLatch latch((int)target_size);

        for (int i = 0; i < target_size; i++) {

            submit_status = thread_pool->submit_func([&, i]() {

                SCOPED_ATTACH_TASK(_query_ctx.get());

                prepare_status[i] = _build_pipeline_tasks_for_instance(i, pipeline_id_to_profile);

                latch.count_down();

            });

            if (LIKELY(submit_status.ok())) {

                submitted_tasks++;

            } else {

                break;

            }

        }

        latch.arrive_and_wait(target_size - submitted_tasks);

        if (UNLIKELY(!submit_status.ok())) {

            return submit_status;

        }

        for (int i = 0; i < submitted_tasks; i++) {

            if (!prepare_status[i].ok()) {

                return prepare_status[i];

            }

        }

    } else {

        for (int i = 0; i < target_size; i++) {

            RETURN_IF_ERROR(_build_pipeline_tasks_for_instance(i, pipeline_id_to_profile));

        }

    }

    _pipeline_parent_map.clear();

    _op_id_to_shared_state.clear();

    // Record task cardinality once when this fragment context finishes task initialization.

    _query_ctx->add_total_task_num(_total_tasks.load(std::memory_order_relaxed));

    return Status::OK();

}

void PipelineFragmentContext::_init_next_report_time() {

    auto interval_s = config::pipeline_status_report_interval;

    if (_is_report_success && interval_s > 0 && _timeout > interval_s) {

        VLOG_FILE << "enable period report: fragment id=" << _fragment_id;

        uint64_t report_fragment_offset = (uint64_t)(rand() % interval_s) * NANOS_PER_SEC;

        // We don't want to wait longer than it takes to run the entire fragment.

        _previous_report_time =

                MonotonicNanos() + report_fragment_offset - (uint64_t)(interval_s)*NANOS_PER_SEC;

        _disable_period_report = false;

    }

}

void PipelineFragmentContext::refresh_next_report_time() {

    auto disable = _disable_period_report.load(std::memory_order_acquire);

    DCHECK(disable == true);

    _previous_report_time.store(MonotonicNanos(), std::memory_order_release);

    _disable_period_report.compare_exchange_strong(disable, false);

}

void PipelineFragmentContext::trigger_report_if_necessary() {

    if (!_is_report_success) {

        return;

    }

    auto disable = _disable_period_report.load(std::memory_order_acquire);

    if (disable) {

        return;

    }

    int32_t interval_s = config::pipeline_status_report_interval;

    if (interval_s <= 0) {

        LOG(WARNING) << "config::status_report_interval is equal to or less than zero, do not "

                        "trigger "

                        "report.";

    }

    uint64_t next_report_time = _previous_report_time.load(std::memory_order_acquire) +

                                (uint64_t)(interval_s)*NANOS_PER_SEC;

    if (MonotonicNanos() > next_report_time) {

        if (!_disable_period_report.compare_exchange_strong(disable, true,

                                                            std::memory_order_acq_rel)) {

            return;

        }

        if (VLOG_FILE_IS_ON) {

            VLOG_FILE << "Reporting "

                      << "profile for query_id " << print_id(_query_id)

                      << ", fragment id: " << _fragment_id;

            std::stringstream ss;

            _runtime_state->runtime_profile()->compute_time_in_profile();

            _runtime_state->runtime_profile()->pretty_print(&ss);

            if (_runtime_state->load_channel_profile()) {

                _runtime_state->load_channel_profile()->pretty_print(&ss);

            }

            VLOG_FILE << "Query " << print_id(get_query_id()) << " fragment " << get_fragment_id()

                      << " profile:\n"

                      << ss.str();

        }

        auto st = send_report(false);

        if (!st.ok()) {

            disable = true;

            _disable_period_report.compare_exchange_strong(disable, false,

                                                           std::memory_order_acq_rel);

        }

    }

}

Status PipelineFragmentContext::_build_pipelines(ObjectPool* pool, const DescriptorTbl& descs,

                                                 OperatorPtr* root, PipelinePtr cur_pipe) {

    if (_params.fragment.plan.nodes.empty()) {

        throw Exception(ErrorCode::INTERNAL_ERROR, "Invalid plan which has no plan node!");

    }

    int node_idx = 0;

    RETURN_IF_ERROR(_create_tree_helper(pool, _params.fragment.plan.nodes, descs, nullptr,

                                        &node_idx, root, cur_pipe, 0, false, false));

    if (node_idx + 1 != _params.fragment.plan.nodes.size()) {

        return Status::InternalError(

                "Plan tree only partially reconstructed. Not all thrift nodes were used.");

    }

    return Status::OK();

}

Status PipelineFragmentContext::_create_deferred_local_exchangers() {

    for (auto& info : _deferred_exchangers) {

        // DANGER ZONE — do not "fix" this line without reading the history.

        //

        // sender_count seeds Exchanger::_running_sink_operators, which the source side

        // waits to reach 0 via sub_running_sink_operators on each sink LocalState close.

        // The correct value is THIS pipeline-instance's sink task count, which is exactly

        // info.upstream_pipe->num_tasks() — one PipelineTask per task, one close per task.

        //

        // Tempting wrong fix #1: `std::max(num_tasks, _num_instances)` to mirror the

        //   BE-planned path in _add_local_exchange_impl (~line 1023).  THIS BREAKS the

        //   common FE-planned shape of `serial scan → LE(PT) → ...`: upstream_pipe

        //   genuinely has num_tasks=1, only 1 close arrives, but seed becomes

        //   _num_instances so _running_sink_operators never reaches 0 — downstream

        //   sources hang on SHUFFLE_DATA_DEPENDENCY (e.g. MTMV refresh from

        //   mtmv_up_down_job_p0/load.groovy stays at Status=RUNNING and regressed

        //   exactly this way).  BE-planned mode uses max() because its

        //   `cur_pipe` is the source-side pipeline (always raised to _num_instances by

        //   add_pipeline) — not analogous to our `upstream_pipe` here, which is the

        //   sink-side pipeline that may legitimately stay at 1 for serial sources.

        //

        // Tempting wrong fix #2: multiply by _num_instances on the theory shared_state

        //   is shared across all instances.  Same hang — each fragment-instance

        //   PipelineFragmentContext has its OWN _op_id_to_shared_state map, so the

        //   exchanger is per-instance, not per-BE.  num_tasks() is already the right

        //   close-count for one instance.

        //

        // If a hang shows up with `_running_sink_operators < 0`, the bug is upstream:

        // _propagate_local_exchange_num_tasks left num_tasks too low (or too high) for

        // this fragment shape.  Fix THAT pass, not this seed value.

        const int sender_count = info.upstream_pipe->num_tasks();

        switch (info.partition_type) {

        case TLocalPartitionType::LOCAL_EXECUTION_HASH_SHUFFLE:

        case TLocalPartitionType::GLOBAL_EXECUTION_HASH_SHUFFLE:

            info.shared_state->exchanger = ShuffleExchanger::create_unique(

                    sender_count, _num_instances, info.num_partitions, info.free_blocks_limit,

                    info.partition_type);

            break;

        case TLocalPartitionType::BUCKET_HASH_SHUFFLE:

            info.shared_state->exchanger = BucketShuffleExchanger::create_unique(

                    sender_count, _num_instances, info.num_partitions, info.free_blocks_limit);

            break;

        case TLocalPartitionType::PASSTHROUGH:

            info.shared_state->exchanger = PassthroughExchanger::create_unique(

                    sender_count, _num_instances, info.free_blocks_limit);

            break;

        case TLocalPartitionType::BROADCAST:

            info.shared_state->exchanger = BroadcastExchanger::create_unique(

                    sender_count, _num_instances, info.free_blocks_limit);

            break;

        case TLocalPartitionType::PASS_TO_ONE:

            if (_runtime_state->enable_share_hash_table_for_broadcast_join()) {

                info.shared_state->exchanger = PassToOneExchanger::create_unique(

                        sender_count, _num_instances, info.free_blocks_limit);

            } else {

                info.shared_state->exchanger = BroadcastExchanger::create_unique(

                        sender_count, _num_instances, info.free_blocks_limit);

            }

            break;

        case TLocalPartitionType::ADAPTIVE_PASSTHROUGH:

            info.shared_state->exchanger = AdaptivePassthroughExchanger::create_unique(

                    sender_count, _num_instances, info.free_blocks_limit);

            break;

        case TLocalPartitionType::NOOP:

        case TLocalPartitionType::LOCAL_MERGE_SORT:

            // FE-planned LocalExchangeNode currently never emits NOOP or LOCAL_MERGE_SORT

            // through the deferred-exchanger path.  NOOP means "no exchange needed" and

            // is filtered out before reaching here; LOCAL_MERGE_SORT is planned by the

            // legacy BE path only.  Crash in debug to surface the protocol violation if

            // that ever changes; return an error in release to avoid silently corrupting

            // execution.

            DCHECK(false) << "FE-planned local exchange should not emit partition_type="

                          << static_cast<int>(info.partition_type);

            return Status::InternalError("FE-planned local exchange emitted unsupported type: " +

                                         std::to_string(static_cast<int>(info.partition_type)));

        default:

            // New TLocalPartitionType added on FE side without a BE handler here.

            DCHECK(false) << "Unhandled TLocalPartitionType in deferred exchangers: "

                          << static_cast<int>(info.partition_type);

            return Status::InternalError("Unsupported FE-planned local exchange type: " +

                                         std::to_string(static_cast<int>(info.partition_type)));

        }

    }

    _deferred_exchangers.clear();

    return Status::OK();

}

void PipelineFragmentContext::_propagate_local_exchange_num_tasks() {

    // Only runs when FE has planned local exchanges and BE deferred their construction.

    // In legacy mode (enable_local_shuffle_planner=false) BE plans LE itself via

    // _plan_local_exchange and _deferred_exchangers stays empty — the legacy path

    // already gets its num_tasks right at construction time, so the propagate passes

    // would be no-ops and are skipped.  This is a transitional design: once the FE

    // planner is the only planner, the propagation logic itself should degrade into

    // a pure assertion that the FE plan already wired the right num_tasks everywhere.

    if (_deferred_exchangers.empty()) {

        return;

    }

    // Reconcile num_tasks across paired pipelines created by pipeline-splitting operators

    // (AGG, SORT, JOIN): they share state via inject_shared_state and must agree, or

    // instance 1+ tasks access null shared_state.  A pipeline's num_tasks is fully

    // determined by its source operator plus its upstreams:

    //   - LocalExchangeSource  -> _num_instances (the LE re-parallelizes)

    //   - serial source        -> its reduced count (kept as-is, typically 1)

    //   - otherwise (splitter) -> inherit from upstreams: raise to _num_instances if any

    //                             upstream was raised by an LE, then lower to a serial

    //                             upstream's count (lower wins).

    // Visiting each pipeline only after all its upstreams (topological order over _dag) lets

    // a single sweep reach the same fixpoint the previous two while-loops iterated to — those

    // only existed to reconcile the top-down build's parent-inherited num_tasks guesses.

    std::map<PipelineId, PipelinePtr> id_to_pipe;

    std::map<PipelineId, std::vector<PipelineId>> downstreams_of;

    std::map<PipelineId, int> in_degree;

    for (auto& p : _pipelines) {

        id_to_pipe[p->id()] = p;

        in_degree.try_emplace(p->id(), 0);

    }

    for (const auto& [downstream_id, upstream_ids] : _dag) {

        for (auto upstream_id : upstream_ids) {

            downstreams_of[upstream_id].push_back(downstream_id);

            in_degree[downstream_id]++;

        }

    }

    std::vector<PipelineId> ready;

    for (const auto& [id, deg] : in_degree) {

        if (deg == 0) {

            ready.push_back(id);

        }

    }

    size_t visited = 0;

    while (!ready.empty()) {

        const auto id = ready.back();

        ready.pop_back();

        visited++;

        auto pit = id_to_pipe.find(id);

        if (pit != id_to_pipe.end()) {

            auto& pipe = pit->second;

            const auto& ops = pipe->operators();

            const bool le_source =

                    !ops.empty() && dynamic_cast<LocalExchangeSourceOperatorX*>(ops.front().get());

            const bool serial_source = !ops.empty() && ops.front()->is_serial_operator();

            if (le_source) {

                pipe->set_num_tasks(_num_instances);

            } else if (!serial_source) {

                int target = pipe->num_tasks();

                const auto up_it = _dag.find(id);

                if (up_it != _dag.end()) {

                    // raise: any upstream already at _num_instances (e.g. an LE source)

                    for (auto upstream_id : up_it->second) {

                        auto uit = id_to_pipe.find(upstream_id);

                        if (uit != id_to_pipe.end() && uit->second->num_tasks() >= _num_instances) {

                            target = _num_instances;

                            break;

                        }

                    }

                    // lower: a serial upstream with fewer tasks (wins over the raise above)

                    for (auto upstream_id : up_it->second) {

                        auto uit = id_to_pipe.find(upstream_id);

                        if (uit != id_to_pipe.end() && uit->second->num_tasks() < target &&

                            !uit->second->operators().empty() &&

                            uit->second->operators().front()->is_serial_operator()) {

                            target = uit->second->num_tasks();

                        }

                    }

                }

                pipe->set_num_tasks(target);

            }

        }

        for (auto down : downstreams_of[id]) {

            if (--in_degree[down] == 0) {

                ready.push_back(down);

            }

        }

    }

    // The pipeline DAG is acyclic; if a future change introduces a back-edge, some pipelines

    // stay unvisited (in_degree never reaches 0) — fail loudly rather than silently leaving

    // their num_tasks unreconciled.

    DCHECK_EQ(visited, in_degree.size())

            << "pipeline num_tasks topological sweep visited " << visited << " of "

            << in_degree.size() << " pipelines (cycle in _dag?)";

}

Status PipelineFragmentContext::_create_tree_helper(

        ObjectPool* pool, const std::vector<TPlanNode>& tnodes, const DescriptorTbl& descs,

        OperatorPtr parent, int* node_idx, OperatorPtr* root, PipelinePtr& cur_pipe, int child_idx,

        const bool followed_by_shuffled_operator, const bool require_bucket_distribution) {

    // propagate error case

    if (*node_idx >= tnodes.size()) {

        return Status::InternalError(

                "Failed to reconstruct plan tree from thrift. Node id: {}, number of nodes: {}",

                *node_idx, tnodes.size());

    }

    const TPlanNode& tnode = tnodes[*node_idx];

    int num_children = tnodes[*node_idx].num_children;

    bool current_followed_by_shuffled_operator = followed_by_shuffled_operator;

    bool current_require_bucket_distribution = require_bucket_distribution;

    // TODO: Create CacheOperator is confused now

    OperatorPtr op = nullptr;

    OperatorPtr cache_op = nullptr;

    RETURN_IF_ERROR(_create_operator(pool, tnodes[*node_idx], descs, op, cur_pipe,

                                     parent == nullptr ? -1 : parent->node_id(), child_idx,

                                     followed_by_shuffled_operator,

                                     current_require_bucket_distribution, cache_op));

    // Initialization must be done here. For example, group by expressions in agg will be used to

    // decide if a local shuffle should be planed, so it must be initialized here.

    RETURN_IF_ERROR(op->init(tnode, _runtime_state.get()));

    // assert(parent != nullptr || (node_idx == 0 && root_expr != nullptr));

    if (parent != nullptr) {

        // add to parent's child(s)

        RETURN_IF_ERROR(parent->set_child(cache_op ? cache_op : op));

    } else {

        *root = op;

    }

    /**

     * `TLocalPartitionType::GLOBAL_EXECUTION_HASH_SHUFFLE` should be used if an operator is followed by a shuffled operator (shuffled hash join, union operator followed by co-located operators).

     *

     * For plan:

     * LocalExchange(id=0) -> Aggregation(id=1) -> ShuffledHashJoin(id=2)

     *                           Exchange(id=3) -> ShuffledHashJoinBuild(id=2)

     * We must ensure data distribution of `LocalExchange(id=0)` is same as Exchange(id=3).

     *

     * If an operator's is followed by a local exchange without shuffle (e.g. passthrough), a

     * shuffled local exchanger will be used before join so it is not followed by shuffle join.

     */

    auto required_data_distribution =

            cur_pipe->operators().empty()

                    ? cur_pipe->sink()->required_data_distribution(_runtime_state.get())

                    : op->required_data_distribution(_runtime_state.get());

    current_followed_by_shuffled_operator =

            ((followed_by_shuffled_operator ||

              (cur_pipe->operators().empty() ? cur_pipe->sink()->is_shuffled_operator()

                                             : op->is_shuffled_operator())) &&

             Pipeline::is_hash_exchange(required_data_distribution.distribution_type)) ||

            (followed_by_shuffled_operator &&

             required_data_distribution.distribution_type == TLocalPartitionType::NOOP);

    current_require_bucket_distribution =

            ((require_bucket_distribution ||

              (cur_pipe->operators().empty() ? cur_pipe->sink()->is_colocated_operator()

                                             : op->is_colocated_operator())) &&

             Pipeline::is_hash_exchange(required_data_distribution.distribution_type)) ||

            (require_bucket_distribution &&

             required_data_distribution.distribution_type == TLocalPartitionType::NOOP);

    if (num_children == 0) {

        _use_serial_source = op->is_serial_operator();

    }

    // rely on that tnodes is preorder of the plan

    for (int i = 0; i < num_children; i++) {

        ++*node_idx;

        RETURN_IF_ERROR(_create_tree_helper(pool, tnodes, descs, op, node_idx, nullptr, cur_pipe, i,

                                            current_followed_by_shuffled_operator,

                                            current_require_bucket_distribution));

        // we are expecting a child, but have used all nodes

        // this means we have been given a bad tree and must fail

        if (*node_idx >= tnodes.size()) {

            return Status::InternalError(

                    "Failed to reconstruct plan tree from thrift. Node id: {}, number of "

                    "nodes: {}",

                    *node_idx, tnodes.size());

        }

    }

    return Status::OK();

}

void PipelineFragmentContext::_inherit_pipeline_properties(

        const DataDistribution& data_distribution, PipelinePtr pipe_with_source,

        PipelinePtr pipe_with_sink) {

    pipe_with_sink->set_num_tasks(pipe_with_source->num_tasks());

    pipe_with_source->set_num_tasks(_num_instances);

    pipe_with_source->set_data_distribution(data_distribution);

}

Status PipelineFragmentContext::_add_local_exchange_impl(

        int idx, ObjectPool* pool, PipelinePtr cur_pipe, PipelinePtr new_pip,

        DataDistribution data_distribution, bool* do_local_exchange, int num_buckets,

        const std::map<int, int>& bucket_seq_to_instance_idx,

        const std::map<int, int>& shuffle_idx_to_instance_idx) {

    auto& operators = cur_pipe->operators();

    const auto downstream_pipeline_id = cur_pipe->id();

    auto local_exchange_id = next_operator_id();

    // 1. Create a new pipeline with local exchange sink.

    DataSinkOperatorPtr sink;

    auto sink_id = next_sink_operator_id();

    /**

     * `bucket_seq_to_instance_idx` is empty if no scan operator is contained in this fragment.

     * So co-located operators(e.g. Agg, Analytic) should use `HASH_SHUFFLE` instead of `BUCKET_HASH_SHUFFLE`.

     */

    const bool followed_by_shuffled_operator =

            operators.size() > idx ? operators[idx]->followed_by_shuffled_operator()

                                   : cur_pipe->sink()->followed_by_shuffled_operator();

    const bool use_global_hash_shuffle = bucket_seq_to_instance_idx.empty() &&

                                         !shuffle_idx_to_instance_idx.contains(-1) &&

                                         followed_by_shuffled_operator && !_use_serial_source;

    sink = std::make_shared<LocalExchangeSinkOperatorX>(

            sink_id, local_exchange_id, use_global_hash_shuffle ? _total_instances : _num_instances,

            data_distribution.partition_exprs, bucket_seq_to_instance_idx);

    if (bucket_seq_to_instance_idx.empty() &&

        data_distribution.distribution_type == TLocalPartitionType::BUCKET_HASH_SHUFFLE) {

        data_distribution.distribution_type =

                use_global_hash_shuffle ? TLocalPartitionType::GLOBAL_EXECUTION_HASH_SHUFFLE

                                        : TLocalPartitionType::LOCAL_EXECUTION_HASH_SHUFFLE;

    }

    if (!use_global_hash_shuffle &&

        data_distribution.distribution_type == TLocalPartitionType::GLOBAL_EXECUTION_HASH_SHUFFLE) {

        data_distribution.distribution_type = TLocalPartitionType::LOCAL_EXECUTION_HASH_SHUFFLE;

    }