// 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/sink/vrow_distribution.h"

#include <gen_cpp/FrontendService.h>

#include <gen_cpp/FrontendService_types.h>

#include <glog/logging.h>

#include <cstdint>

#include <memory>

#include <optional>

#include <string>

#include "common/cast_set.h"

#include "common/logging.h"

#include "common/status.h"

#include "core/assert_cast.h"

#include "core/column/column.h"

#include "core/column/column_const.h"

#include "core/column/column_nullable.h"

#include "core/column/column_vector.h"

#include "core/data_type/data_type.h"

#include "exec/sink/writer/vtablet_writer.h"

#include "runtime/exec_env.h"

#include "runtime/query_context.h"

#include "runtime/runtime_state.h"

#include "service/backend_options.h"

#include "util/client_cache.h"

#include "util/debug_points.h"

#include "util/thrift_rpc_helper.h"

namespace doris {

std::pair<VExprContextSPtrs, VExprSPtrs> VRowDistribution::_get_partition_function() {

    return {_vpartition->get_part_func_ctx(), _vpartition->get_partition_function()};

}

Status VRowDistribution::_save_missing_values(

        const Block& input_block,

        std::vector<std::vector<std::string>>& col_strs, // non-const ref for move

        int col_size, Block* block, const std::vector<uint32_t>& filter,

        const std::vector<std::optional<NullMapView>>& col_null_maps) {

    // de-duplication for new partitions but save all rows.

    RETURN_IF_ERROR(

            _batching_block->add_rows(&input_block, filter.data(), filter.data() + filter.size()));

    std::vector<TNullableStringLiteral> cur_row_values;

    for (int row = 0; row < col_strs[0].size(); ++row) {

        cur_row_values.clear();

        for (int col = 0; col < col_size; ++col) {

            TNullableStringLiteral node;

            // OlapTableBlockConvertor::_validate_data() materializes destination slots so won't be const.

            const auto& null_map = col_null_maps[col]; // null map for this col

            node.__set_is_null((null_map.has_value() && (*null_map)[filter[row]])

                                       ? true

                                       : node.is_null); // if not, dont change(default false)

            if (!node.is_null) {

                node.__set_value(col_strs[col][row]);

            }

            cur_row_values.push_back(node);

        }

        if (!_deduper.contains(cur_row_values)) {

            _deduper.insert(cur_row_values);

            _partitions_need_create.emplace_back(cur_row_values);

        }

    }

    // to avoid too large mem use

    if (_batching_block->rows() > _batch_size) {

        _deal_batched = true;

    }

    _batching_rows = _batching_block->rows();

    VLOG_NOTICE << "pushed some batching lines, now numbers = " << _batching_rows;

    return Status::OK();

}

void VRowDistribution::clear_batching_stats() {

    _partitions_need_create.clear();

    _batching_rows = 0;

    _batching_bytes = 0;

}

Status VRowDistribution::automatic_create_partition() {

    MonotonicStopWatch timer;

    if (_state->enable_profile() && _state->profile_level() >= 2) {

        timer.start();

    }

    SCOPED_TIMER(_add_partition_request_timer);

    TCreatePartitionRequest request;

    TCreatePartitionResult result;

    bool injected = false;

    std::string be_endpoint = BackendOptions::get_be_endpoint();

    request.__set_txn_id(_txn_id);

    request.__set_db_id(_vpartition->db_id());

    request.__set_table_id(_vpartition->table_id());

    request.__set_partitionValues(_partitions_need_create);

    request.__set_be_endpoint(be_endpoint);

    request.__set_write_single_replica(_write_single_replica);

    request.__set_load_to_single_tablet(_tablet_finder->is_find_tablet_every_sink());

    if (_state && _state->get_query_ctx()) {

        // Pass query_id to FE so it can determine if this is a multi-instance load by checking Coordinator

        request.__set_query_id(_state->get_query_ctx()->query_id());

    }

    DBUG_EXECUTE_IF("VRowDistribution.automatic_create_partition.inject_result", {

        DBUG_RUN_CALLBACK(&request, &result);

        injected = true;

    });

    VLOG_NOTICE << "automatic partition rpc begin request " << request;

    if (!injected) {

        std::shared_ptr<TNetworkAddress> master_addr;

        if (_vpartition->get_master_address() == nullptr) {

            auto* cluster_info = ExecEnv::GetInstance()->cluster_info();

            if (cluster_info == nullptr) {

                return Status::InternalError("cluster_info is null");

            }

            master_addr = std::make_shared<TNetworkAddress>(cluster_info->master_fe_addr);

        } else {

            master_addr = _vpartition->get_master_address();

        }

        int time_out = _state->execution_timeout() * 1000;

        RETURN_IF_ERROR(ThriftRpcHelper::rpc<FrontendServiceClient>(

                master_addr->hostname, master_addr->port,

                [&request, &result](FrontendServiceConnection& client) {

                    client->createPartition(result, request);

                },

                time_out));

    }

    Status status(Status::create(result.status));

    VLOG_NOTICE << "automatic partition rpc end response " << result;

    if (result.status.status_code == TStatusCode::OK) {

        RETURN_IF_ERROR(_create_partition_callback(_caller, &result));

        // add new created partitions

        RETURN_IF_ERROR(_vpartition->add_partitions(result.partitions));

        for (const auto& part : result.partitions) {

            _new_partition_ids.insert(part.id);

            VLOG_TRACE << "record new id: " << part.id;

        }

    }

    // Record this request's elapsed time

    if (_state->enable_profile() && _state->profile_level() >= 2) {

        int64_t elapsed_ns = timer.elapsed_time();

        _add_partition_request_times.push_back(elapsed_ns);

    }

    return status;

}

// for reuse the same create callback of create-partition

static TCreatePartitionResult cast_as_create_result(const TReplacePartitionResult& arg) {

    TCreatePartitionResult result;

    result.status = arg.status;

    result.nodes = arg.nodes;

    result.partitions = arg.partitions;

    result.tablets = arg.tablets;

    result.slave_tablets = arg.slave_tablets;

    return result;

}

// use _partitions and replace them

Status VRowDistribution::_replace_overwriting_partition() {

    SCOPED_TIMER(_add_partition_request_timer); // also for replace_partition

    TReplacePartitionRequest request;

    TReplacePartitionResult result;

    bool injected = false;

    request.__set_overwrite_group_id(_vpartition->get_overwrite_group_id());

    request.__set_db_id(_vpartition->db_id());

    request.__set_table_id(_vpartition->table_id());

    request.__set_write_single_replica(_write_single_replica);

    // only request for partitions not recorded for replacement

    std::set<int64_t> id_deduper;

    for (const auto* part : _partitions) {

        if (part != nullptr) {

            if (_new_partition_ids.contains(part->id)) {

                // this is a new partition. dont replace again.

                VLOG_TRACE << "skip new partition: " << part->id;

            } else {

                // request for replacement

                id_deduper.insert(part->id);

            }

        } else if (_missing_map.empty()) {

            // no origin partition. and not allow to create.

            return Status::InvalidArgument(

                    "Cannot found origin partitions in auto detect overwriting, stop "

                    "processing");

        } // else: part is null and _missing_map is not empty. dealed outside using auto-partition way. nothing to do here.

    }

    if (id_deduper.empty()) {

        return Status::OK(); // no need to request

    }

    // de-duplicate. there's no check in FE

    std::vector<int64_t> request_part_ids(id_deduper.begin(), id_deduper.end());

    request.__set_partition_ids(request_part_ids);

    std::string be_endpoint = BackendOptions::get_be_endpoint();

    request.__set_be_endpoint(be_endpoint);

    request.__set_load_to_single_tablet(_tablet_finder->is_find_tablet_every_sink());

    if (_state && _state->get_query_ctx()) {

        // Pass query_id to FE so it can determine if this is a multi-instance load by checking Coordinator

        request.__set_query_id(_state->get_query_ctx()->query_id());

    }

    DBUG_EXECUTE_IF("VRowDistribution.replace_overwriting_partition.inject_result", {

        DBUG_RUN_CALLBACK(&request, &result);

        injected = true;

    });

    VLOG_NOTICE << "auto detect replace partition request: " << request;

    if (!injected) {

        std::shared_ptr<TNetworkAddress> master_addr;

        if (_vpartition->get_master_address() == nullptr) {

            auto* cluster_info = ExecEnv::GetInstance()->cluster_info();

            if (cluster_info == nullptr) {

                return Status::InternalError("cluster_info is null");

            }

            master_addr = std::make_shared<TNetworkAddress>(cluster_info->master_fe_addr);

        } else {

            master_addr = _vpartition->get_master_address();

        }

        int time_out = _state->execution_timeout() * 1000;

        RETURN_IF_ERROR(ThriftRpcHelper::rpc<FrontendServiceClient>(

                master_addr->hostname, master_addr->port,

                [&request, &result](FrontendServiceConnection& client) {

                    client->replacePartition(result, request);

                },

                time_out));

    }

    Status status(Status::create(result.status));

    VLOG_NOTICE << "auto detect replace partition result: " << result;

    if (result.status.status_code == TStatusCode::OK) {

        // Reuse the function as the args' structure are same. It adds nodes/locations

        // and waits for incremental_open before the new tablets become routable.

        auto result_as_create = cast_as_create_result(result);

        RETURN_IF_ERROR(_create_partition_callback(_caller, &result_as_create));

        // record new partitions

        for (const auto& part : result.partitions) {

            _new_partition_ids.insert(part.id);

            VLOG_TRACE << "record new id: " << part.id;

        }

        // replace data in _partitions

        RETURN_IF_ERROR(_vpartition->replace_partitions(request_part_ids, result.partitions));

    }

    return status;

}

void VRowDistribution::_get_tablet_ids(Block* block, int32_t index_idx,

                                       std::vector<int64_t>& tablet_ids) {

    tablet_ids.resize(block->rows());

    for (int row_idx = 0; row_idx < block->rows(); row_idx++) {

        if (_skip[row_idx]) {

            continue;

        }

        auto& partition = _partitions[row_idx];

        auto& tablet_index = _tablet_indexes[row_idx];

        auto& index = partition->indexes[index_idx];

        auto tablet_id = index.tablets[tablet_index];

        tablet_ids[row_idx] = tablet_id;

    }

}

void VRowDistribution::_filter_block_by_skip(Block* block, RowPartTabletIds& row_part_tablet_id) {

    auto& row_ids = row_part_tablet_id.row_ids;

    auto& partition_ids = row_part_tablet_id.partition_ids;

    auto& tablet_ids = row_part_tablet_id.tablet_ids;

    auto rows = block->rows();

    // row count of a block should not exceed UINT32_MAX

    auto rows_uint32 = cast_set<uint32_t>(rows);

    for (uint32_t i = 0; i < rows_uint32; i++) {

        if (!_skip[i]) {

            row_ids.emplace_back(i);

            partition_ids.emplace_back(_partitions[i]->id);

            tablet_ids.emplace_back(_tablet_ids[i]);

        }

    }

}

Status VRowDistribution::_filter_block_by_skip_and_where_clause(

        Block* block, const VExprContextSPtr& where_clause, RowPartTabletIds& row_part_tablet_id) {

    // TODO

    //SCOPED_RAW_TIMER(&_stat.where_clause_ns);

    ColumnPtr filter_column;

    RETURN_IF_ERROR(where_clause->execute(block, filter_column));

    auto& row_ids = row_part_tablet_id.row_ids;

    auto& partition_ids = row_part_tablet_id.partition_ids;

    auto& tablet_ids = row_part_tablet_id.tablet_ids;

    if (const auto* nullable_column = check_and_get_column<ColumnNullable>(*filter_column)) {

        auto rows = block->rows();

        // row count of a block should not exceed UINT32_MAX

        auto rows_uint32 = cast_set<uint32_t>(rows);

        for (uint32_t i = 0; i < rows_uint32; i++) {

            if (nullable_column->get_bool_inline(i) && !_skip[i]) {

                row_ids.emplace_back(i);

                partition_ids.emplace_back(_partitions[i]->id);

                tablet_ids.emplace_back(_tablet_ids[i]);

            }

        }

    } else if (const auto* const_column = check_and_get_column<ColumnConst>(*filter_column)) {

        bool ret = const_column->get_bool(0);

        if (!ret) {

            return Status::OK();

        }

        // should we optimize?

        _filter_block_by_skip(block, row_part_tablet_id);

    } else {

        const auto& filter = assert_cast<const ColumnUInt8&>(*filter_column).get_data();

        auto rows = block->rows();

        // row count of a block should not exceed UINT32_MAX

        auto rows_uint32 = cast_set<uint32_t>(rows);

        for (uint32_t i = 0; i < rows_uint32; i++) {

            if (filter[i] != 0 && !_skip[i]) {

                row_ids.emplace_back(i);

                partition_ids.emplace_back(_partitions[i]->id);

                tablet_ids.emplace_back(_tablet_ids[i]);

            }

        }

    }

    return Status::OK();

}

Status VRowDistribution::_filter_block(Block* block,

                                       std::vector<RowPartTabletIds>& row_part_tablet_ids) {

    for (int i = 0; i < _schema->indexes().size(); i++) {

        _get_tablet_ids(block, i, _tablet_ids);

        auto& where_clause = _schema->indexes()[i]->where_clause;

        if (where_clause != nullptr) {

            RETURN_IF_ERROR(_filter_block_by_skip_and_where_clause(block, where_clause,

                                                                   row_part_tablet_ids[i]));

        } else {

            _filter_block_by_skip(block, row_part_tablet_ids[i]);

        }

    }

    return Status::OK();

}

Status VRowDistribution::_generate_rows_distribution_for_non_auto_partition(

        Block* block, bool has_filtered_rows, std::vector<RowPartTabletIds>& row_part_tablet_ids) {

    int num_rows = cast_set<int>(block->rows());

    RETURN_IF_ERROR(_tablet_finder->find_tablets(_state, block, num_rows, _partitions,

                                                 _tablet_indexes, _skip));

    if (has_filtered_rows) {

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

            _skip[i] = _skip[i] || _block_convertor->filter_map()[i];

        }

    }

    RETURN_IF_ERROR(_filter_block(block, row_part_tablet_ids));

    return Status::OK();

}

Status VRowDistribution::_deal_missing_map(const Block& input_block, Block* block,

                                           const std::vector<uint16_t>& partition_cols_idx,

                                           int64_t& rows_stat_val) {

    // for missing partition keys, calc the missing partition and save in _partitions_need_create

    auto [part_ctxs, part_exprs] = _get_partition_function();

    int part_col_num = cast_set<int>(part_exprs.size());

    // the two vectors are in column-first-order

    std::vector<std::vector<std::string>> col_strs;

    std::vector<std::optional<NullMapView>> col_null_maps;

    col_strs.resize(part_col_num);

    col_null_maps.reserve(part_col_num);

    auto format_options = DataTypeSerDe::get_default_format_options();

    format_options.timezone = &_state->timezone_obj();

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

        auto return_type = part_exprs[i]->data_type();

        // expose the data column. the return type would be nullable

        const auto& [range_left_col, col_const] =

                unpack_if_const(block->get_by_position(partition_cols_idx[i]).column);

        if (range_left_col->is_nullable()) {

            col_null_maps.emplace_back(

                    assert_cast<const ColumnNullable*>(range_left_col.get())->get_null_map_data());

        } else {

            col_null_maps.emplace_back(std::nullopt);

        }

        for (auto row : _missing_map) {

            col_strs[i].push_back(return_type->to_string(

                    *range_left_col, index_check_const(row, col_const), format_options));

        }

    }

    // calc the end value and save them. in the end of sending, we will create partitions for them and deal them.

    // NOTE: must save old batching stats before calling _save_missing_values(),

    // because _save_missing_values() will update _batching_rows internally.

    size_t old_bt_rows = _batching_rows;

    size_t old_bt_bytes = _batching_bytes;

    RETURN_IF_ERROR(_save_missing_values(input_block, col_strs, part_col_num, block, _missing_map,

                                         col_null_maps));

    size_t new_bt_rows = _batching_block->rows();

    size_t new_bt_bytes = _batching_block->bytes();

    rows_stat_val -= new_bt_rows - old_bt_rows;

    _state->update_num_rows_load_total(old_bt_rows - new_bt_rows);

    _state->update_num_bytes_load_total(old_bt_bytes - new_bt_bytes);

    return Status::OK();

}

Status VRowDistribution::_generate_rows_distribution_for_auto_partition(

        const Block& input_block, Block* block, const std::vector<uint16_t>& partition_cols_idx,

        bool has_filtered_rows, std::vector<RowPartTabletIds>& row_part_tablet_ids,

        int64_t& rows_stat_val) {

    int num_rows = cast_set<int>(block->rows());

    std::vector<uint16_t> partition_keys = _vpartition->get_partition_keys();

    auto& partition_col = block->get_by_position(partition_keys[0]);

    _missing_map.clear();

    _missing_map.reserve(partition_col.column->size());

    RETURN_IF_ERROR(_tablet_finder->find_tablets(_state, block, num_rows, _partitions,

                                                 _tablet_indexes, _skip, &_missing_map));

    // the missing vals for auto partition are also skipped.

    if (has_filtered_rows) {

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

            _skip[i] = _skip[i] || _block_convertor->filter_map()[i];

        }

    }

    RETURN_IF_ERROR(_filter_block(block, row_part_tablet_ids));

    if (!_missing_map.empty()) {

        RETURN_IF_ERROR(_deal_missing_map(input_block, block, partition_cols_idx,

                                          rows_stat_val)); // send input block to save

    }

    return Status::OK();

}

Status VRowDistribution::_generate_rows_distribution_for_auto_overwrite(

        const Block& input_block, Block* block, const std::vector<uint16_t>& partition_cols_idx,

        bool has_filtered_rows, std::vector<RowPartTabletIds>& row_part_tablet_ids,

        int64_t& rows_stat_val) {

    int num_rows = cast_set<int>(block->rows());

    // for non-auto-partition situation, goes into two 'else' branch. just find the origin partitions, replace them by rpc,

    //  and find the new partitions to use.

    // for auto-partition's, find and save origins in _partitions and replace them. at meanwhile save the missing values for auto

    //  partition. then we find partition again to get replaced partitions in _partitions. this time _missing_map is ignored cuz

    //  we already saved missing values.

    if (_vpartition->is_auto_partition() &&

        _state->query_options().enable_auto_create_when_overwrite) {

        // allow auto create partition for missing rows.

        std::vector<uint16_t> partition_keys = _vpartition->get_partition_keys();

        auto partition_col = block->get_by_position(partition_keys[0]);

        _missing_map.clear();

        _missing_map.reserve(partition_col.column->size());

        RETURN_IF_ERROR(_tablet_finder->find_tablets(_state, block, num_rows, _partitions,

                                                     _tablet_indexes, _skip, &_missing_map));

        // allow and really need to create during auto-detect-overwriting.

        if (!_missing_map.empty()) {

            RETURN_IF_ERROR(

                    _deal_missing_map(input_block, block, partition_cols_idx, rows_stat_val));

        }

    } else {

        RETURN_IF_ERROR(_tablet_finder->find_tablets(_state, block, num_rows, _partitions,

                                                     _tablet_indexes, _skip));

    }

    RETURN_IF_ERROR(_replace_overwriting_partition());

    // regenerate locations for new partitions & tablets

    _reset_find_tablets(num_rows);

    if (_vpartition->is_auto_partition() &&

        _state->query_options().enable_auto_create_when_overwrite) {

        // here _missing_map is just a placeholder

        RETURN_IF_ERROR(_tablet_finder->find_tablets(_state, block, num_rows, _partitions,

                                                     _tablet_indexes, _skip, &_missing_map));

        if (VLOG_TRACE_IS_ON) {

            std::string tmp;

            for (auto v : _missing_map) {

                tmp += std::to_string(v).append(", ");

            }

            VLOG_TRACE << "Trace missing map of " << this << ':' << tmp;

        }

    } else {

        RETURN_IF_ERROR(_tablet_finder->find_tablets(_state, block, num_rows, _partitions,

                                                     _tablet_indexes, _skip));

    }

    if (has_filtered_rows) {

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

            _skip[i] = _skip[i] || _block_convertor->filter_map()[i];

        }

    }

    RETURN_IF_ERROR(_filter_block(block, row_part_tablet_ids));

    return Status::OK();

}

void VRowDistribution::_reset_row_part_tablet_ids(

        std::vector<RowPartTabletIds>& row_part_tablet_ids, int64_t rows) {

    row_part_tablet_ids.resize(_schema->indexes().size());

    for (auto& row_part_tablet_id : row_part_tablet_ids) {

        auto& row_ids = row_part_tablet_id.row_ids;

        auto& partition_ids = row_part_tablet_id.partition_ids;

        auto& tablet_ids = row_part_tablet_id.tablet_ids;

        row_ids.clear();

        partition_ids.clear();

        tablet_ids.clear();

        // This is important for performance.

        row_ids.reserve(rows);

        partition_ids.reserve(rows);

        tablet_ids.reserve(rows);

    }

}

Status VRowDistribution::generate_rows_distribution(

        Block& input_block, std::shared_ptr<Block>& block,

        std::vector<RowPartTabletIds>& row_part_tablet_ids, int64_t& rows_stat_val) {

    auto input_rows = input_block.rows();

    _reset_row_part_tablet_ids(row_part_tablet_ids, input_rows);

    // we store the batching block with value of `input_block`. so just do all of these again.

    bool has_filtered_rows = false;

    RETURN_IF_ERROR(_block_convertor->validate_and_convert_block(

            _state, &input_block, block, *_vec_output_expr_ctxs, input_rows, has_filtered_rows));

    // batching block rows which need new partitions. deal together at finish.

    if (!_batching_block) [[unlikely]] {

        std::unique_ptr<Block> tmp_block = input_block.create_same_struct_block(0);

        _batching_block = MutableBlock::create_unique(std::move(*tmp_block));

    }

    auto num_rows = block->rows();

    _reset_find_tablets(num_rows);

    // if there's projection of partition calc, we need to calc it first.

    auto [part_ctxs, part_funcs] = _get_partition_function();

    std::vector<uint16_t> partition_cols_idx;

    if (_vpartition->is_projection_partition()) {

        // calc the start value of missing partition ranges.

        auto func_size = part_funcs.size();

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

            int result_idx = -1;

            // we just calc left range here. leave right to FE to avoid dup calc.

            RETURN_IF_ERROR(part_funcs[i]->execute(part_ctxs[i].get(), block.get(), &result_idx));

            VLOG_DEBUG << "Partition-calculated block:\n" << block->dump_data(0, 1);

            DCHECK(result_idx != -1);

            partition_cols_idx.push_back(cast_set<uint16_t>(result_idx));

        }

        // change the column to compare to transformed.

        _vpartition->set_transformed_slots(partition_cols_idx);

    }

    Status st = Status::OK();

    if (_vpartition->is_auto_detect_overwrite() && !_deal_batched) {

        // when overwrite, no auto create partition allowed.

        st = _generate_rows_distribution_for_auto_overwrite(input_block, block.get(),

                                                            partition_cols_idx, has_filtered_rows,

                                                            row_part_tablet_ids, rows_stat_val);

    } else if (_vpartition->is_auto_partition() && !_deal_batched) {

        st = _generate_rows_distribution_for_auto_partition(input_block, block.get(),

                                                            partition_cols_idx, has_filtered_rows,

                                                            row_part_tablet_ids, rows_stat_val);

    } else { // not auto partition

        st = _generate_rows_distribution_for_non_auto_partition(block.get(), has_filtered_rows,

                                                                row_part_tablet_ids);

    }

    return st;

}

// reuse vars for find_tablets

void VRowDistribution::_reset_find_tablets(int64_t rows) {

    _tablet_finder->filter_bitmap().Reset(rows);

    _partitions.assign(rows, nullptr);

    _skip.assign(rows, false);

    _tablet_indexes.assign(rows, 0);

}

} // namespace doris