// 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/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 <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/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/es_scan_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_block_buffer.h"

#include "runtime/result_buffer_mgr.h"

#include "runtime/runtime_state.h"

#include "runtime/thread_context.h"

#include "service/backend_options.h"

#include "util/countdown_latch.h"

#include "util/debug_util.h"

#include "util/uid_util.h"

namespace doris {

#include "common/compile_check_begin.h"

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,

        report_status_callback report_status_cb)

        : _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),

          _report_status_cb(std::move(report_status_cb)),

          _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;

}

// 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());

    {

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

        if (_closed_tasks >= _total_tasks) {

            // All tasks in this PipelineXFragmentContext already closed.

            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->terminate();

        }

    }

}

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));

        // 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->enable_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 (_need_notify_close) {

                    // rec cte require child rf to wait infinitely to make sure all rpc done

                    task_runtime_state->set_force_make_rf_wait_infinite();

                }

                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();

    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_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;

    OperatorPtr 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));

    // 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(op));

    } else {

        *root = op;

    }

    /**

     * `ExchangeType::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 == ExchangeType::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 == ExchangeType::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 == ExchangeType::BUCKET_HASH_SHUFFLE) {

        data_distribution.distribution_type = ExchangeType::HASH_SHUFFLE;

    }

    RETURN_IF_ERROR(new_pip->set_sink(sink));

    RETURN_IF_ERROR(new_pip->sink()->init(_runtime_state.get(), data_distribution.distribution_type,

                                          num_buckets, use_global_hash_shuffle,

                                          shuffle_idx_to_instance_idx));

    // 2. Create and initialize LocalExchangeSharedState.

    std::shared_ptr<LocalExchangeSharedState> shared_state =

            LocalExchangeSharedState::create_shared(_num_instances);

    switch (data_distribution.distribution_type) {

    case ExchangeType::HASH_SHUFFLE:

        shared_state->exchanger = ShuffleExchanger::create_unique(

                std::max(cur_pipe->num_tasks(), _num_instances), _num_instances,

                use_global_hash_shuffle ? _total_instances : _num_instances,

                _runtime_state->query_options().__isset.local_exchange_free_blocks_limit

                        ? cast_set<int>(

                                  _runtime_state->query_options().local_exchange_free_blocks_limit)

                        : 0);

        break;

    case ExchangeType::BUCKET_HASH_SHUFFLE:

        shared_state->exchanger = BucketShuffleExchanger::create_unique(

                std::max(cur_pipe->num_tasks(), _num_instances), _num_instances, num_buckets,

                _runtime_state->query_options().__isset.local_exchange_free_blocks_limit

                        ? cast_set<int>(

                                  _runtime_state->query_options().local_exchange_free_blocks_limit)

                        : 0);

        break;

    case ExchangeType::PASSTHROUGH:

        shared_state->exchanger = PassthroughExchanger::create_unique(

                cur_pipe->num_tasks(), _num_instances,

                _runtime_state->query_options().__isset.local_exchange_free_blocks_limit

                        ? cast_set<int>(

                                  _runtime_state->query_options().local_exchange_free_blocks_limit)

                        : 0);

        break;

    case ExchangeType::BROADCAST:

        shared_state->exchanger = BroadcastExchanger::create_unique(

                cur_pipe->num_tasks(), _num_instances,

                _runtime_state->query_options().__isset.local_exchange_free_blocks_limit

                        ? cast_set<int>(

                                  _runtime_state->query_options().local_exchange_free_blocks_limit)

                        : 0);

        break;

    case ExchangeType::PASS_TO_ONE:

        if (_runtime_state->enable_share_hash_table_for_broadcast_join()) {

            // If shared hash table is enabled for BJ, hash table will be built by only one task

            shared_state->exchanger = PassToOneExchanger::create_unique(

                    cur_pipe->num_tasks(), _num_instances,

                    _runtime_state->query_options().__isset.local_exchange_free_blocks_limit

                            ? cast_set<int>(_runtime_state->query_options()

                                                    .local_exchange_free_blocks_limit)

                            : 0);

        } else {

            shared_state->exchanger = BroadcastExchanger::create_unique(

                    cur_pipe->num_tasks(), _num_instances,

                    _runtime_state->query_options().__isset.local_exchange_free_blocks_limit

                            ? cast_set<int>(_runtime_state->query_options()

                                                    .local_exchange_free_blocks_limit)

                            : 0);

        }

        break;

    case ExchangeType::ADAPTIVE_PASSTHROUGH:

        shared_state->exchanger = AdaptivePassthroughExchanger::create_unique(

                std::max(cur_pipe->num_tasks(), _num_instances), _num_instances,

                _runtime_state->query_options().__isset.local_exchange_free_blocks_limit

                        ? cast_set<int>(

                                  _runtime_state->query_options().local_exchange_free_blocks_limit)

                        : 0);

        break;

    default:

        return Status::InternalError("Unsupported local exchange type : " +

                                     std::to_string((int)data_distribution.distribution_type));

    }

    shared_state->create_source_dependencies(_num_instances, local_exchange_id, local_exchange_id,

                                             "LOCAL_EXCHANGE_OPERATOR");

    shared_state->create_sink_dependency(sink_id, local_exchange_id, "LOCAL_EXCHANGE_SINK");

    _op_id_to_shared_state.insert({local_exchange_id, {shared_state, shared_state->sink_deps}});

    // 3. Set two pipelines' operator list. For example, split pipeline [Scan - AggSink] to

    // pipeline1 [Scan - LocalExchangeSink] and pipeline2 [LocalExchangeSource - AggSink].

    // 3.1 Initialize new pipeline's operator list.

    std::copy(operators.begin(), operators.begin() + idx,

              std::inserter(new_pip->operators(), new_pip->operators().end()));

    // 3.2 Erase unused operators in previous pipeline.

    operators.erase(operators.begin(), operators.begin() + idx);

    // 4. Initialize LocalExchangeSource and insert it into this pipeline.

    OperatorPtr source_op;

    source_op = std::make_shared<LocalExchangeSourceOperatorX>(pool, local_exchange_id);

    RETURN_IF_ERROR(source_op->set_child(new_pip->operators().back()));

    RETURN_IF_ERROR(source_op->init(data_distribution.distribution_type));

    if (!operators.empty()) {

        RETURN_IF_ERROR(operators.front()->set_child(nullptr));

        RETURN_IF_ERROR(operators.front()->set_child(source_op));

    }

    operators.insert(operators.begin(), source_op);

    // 5. Set children for two pipelines separately.

    std::vector<std::shared_ptr<Pipeline>> new_children;

    std::vector<PipelineId> edges_with_source;

    for (auto child : cur_pipe->children()) {

        bool found = false;

        for (auto op : new_pip->operators()) {

            if (child->sink()->node_id() == op->node_id()) {

                new_pip->set_children(child);

                found = true;

            };

        }

        if (!found) {

            new_children.push_back(child);

            edges_with_source.push_back(child->id());

        }

    }

    new_children.push_back(new_pip);

    edges_with_source.push_back(new_pip->id());

    // 6. Set DAG for new pipelines.

    if (!new_pip->children().empty()) {

        std::vector<PipelineId> edges_with_sink;

        for (auto child : new_pip->children()) {

            edges_with_sink.push_back(child->id());

        }

        _dag.insert({new_pip->id(), edges_with_sink});

    }

    cur_pipe->set_children(new_children);

    _dag[downstream_pipeline_id] = edges_with_source;

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

    RETURN_IF_ERROR(cur_pipe->sink()->set_child(nullptr));

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

    // 7. Inherit properties from current pipeline.

    _inherit_pipeline_properties(data_distribution, cur_pipe, new_pip);

    return Status::OK();

}

Status PipelineFragmentContext::_add_local_exchange(

        int pip_idx, int idx, int node_id, ObjectPool* pool, PipelinePtr cur_pipe,

        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) {

    if (_num_instances <= 1 || cur_pipe->num_tasks_of_parent() <= 1) {

        return Status::OK();

    }

    if (!cur_pipe->need_to_local_exchange(data_distribution, idx)) {

        return Status::OK();

    }

    *do_local_exchange = true;

    auto& operators = cur_pipe->operators();

    auto total_op_num = operators.size();

    auto new_pip = add_pipeline(cur_pipe, pip_idx + 1);

    RETURN_IF_ERROR(_add_local_exchange_impl(

            idx, pool, cur_pipe, new_pip, data_distribution, do_local_exchange, num_buckets,

            bucket_seq_to_instance_idx, shuffle_idx_to_instance_idx));

    CHECK(total_op_num + 1 == cur_pipe->operators().size() + new_pip->operators().size())

            << "total_op_num: " << total_op_num

            << " cur_pipe->operators().size(): " << cur_pipe->operators().size()

            << " new_pip->operators().size(): " << new_pip->operators().size();

    // There are some local shuffles with relatively heavy operations on the sink.

    // If the local sink concurrency is 1 and the local source concurrency is n, the sink becomes a bottleneck.

    // Therefore, local passthrough is used to increase the concurrency of the sink.

    // op -> local sink(1) -> local source (n)

    // op -> local passthrough(1) -> local passthrough(n) ->  local sink(n) -> local source (n)

    if (cur_pipe->num_tasks() > 1 && new_pip->num_tasks() == 1 &&

        Pipeline::heavy_operations_on_the_sink(data_distribution.distribution_type)) {

        RETURN_IF_ERROR(_add_local_exchange_impl(

                cast_set<int>(new_pip->operators().size()), pool, new_pip,

                add_pipeline(new_pip, pip_idx + 2), DataDistribution(ExchangeType::PASSTHROUGH),

                do_local_exchange, num_buckets, bucket_seq_to_instance_idx,

                shuffle_idx_to_instance_idx));

    }

    return Status::OK();

}

Status PipelineFragmentContext::_plan_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) {

    for (int pip_idx = cast_set<int>(_pipelines.size()) - 1; pip_idx >= 0; pip_idx--) {

        _pipelines[pip_idx]->init_data_distribution(_runtime_state.get());

        // Set property if child pipeline is not join operator's child.

        if (!_pipelines[pip_idx]->children().empty()) {

            for (auto& child : _pipelines[pip_idx]->children()) {

                if (child->sink()->node_id() ==

                    _pipelines[pip_idx]->operators().front()->node_id()) {

                    _pipelines[pip_idx]->set_data_distribution(child->data_distribution());

                }

            }

        }

        // if 'num_buckets == 0' means the fragment is colocated by exchange node not the

        // scan node. so here use `_num_instance` to replace the `num_buckets` to prevent dividing 0

        // still keep colocate plan after local shuffle

        RETURN_IF_ERROR(_plan_local_exchange(num_buckets, pip_idx, _pipelines[pip_idx],

                                             bucket_seq_to_instance_idx,

                                             shuffle_idx_to_instance_idx));

    }

    return Status::OK();

}

Status PipelineFragmentContext::_plan_local_exchange(

        int num_buckets, int pip_idx, PipelinePtr pip,

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

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

    int idx = 1;

    bool do_local_exchange = false;

    do {

        auto& ops = pip->operators();

        do_local_exchange = false;

        // Plan local exchange for each operator.

        for (; idx < ops.size();) {

            if (ops[idx]->required_data_distribution(_runtime_state.get()).need_local_exchange()) {

                RETURN_IF_ERROR(_add_local_exchange(

                        pip_idx, idx, ops[idx]->node_id(), _runtime_state->obj_pool(), pip,

                        ops[idx]->required_data_distribution(_runtime_state.get()),

                        &do_local_exchange, num_buckets, bucket_seq_to_instance_idx,

                        shuffle_idx_to_instance_idx));

            }

            if (do_local_exchange) {

                // If local exchange is needed for current operator, we will split this pipeline to

                // two pipelines by local exchange sink/source. And then we need to process remaining

                // operators in this pipeline so we set idx to 2 (0 is local exchange source and 1

                // is current operator was already processed) and continue to plan local exchange.

                idx = 2;

                break;

            }

            idx++;

        }

    } while (do_local_exchange);

    if (pip->sink()->required_data_distribution(_runtime_state.get()).need_local_exchange()) {

        RETURN_IF_ERROR(_add_local_exchange(

                pip_idx, idx, pip->sink()->node_id(), _runtime_state->obj_pool(), pip,

                pip->sink()->required_data_distribution(_runtime_state.get()), &do_local_exchange,

                num_buckets, bucket_seq_to_instance_idx, shuffle_idx_to_instance_idx));

    }

    return Status::OK();