// 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 "storage/segment/vertical_segment_writer.h"

#include <crc32c/crc32c.h>

#include <gen_cpp/olap_file.pb.h>

#include <gen_cpp/segment_v2.pb.h>

#include <parallel_hashmap/phmap.h>

#include <cassert>

#include <memory>

#include <ostream>

#include <string>

#include <unordered_map>

#include <unordered_set>

#include <utility>

#include "cloud/config.h"

#include "common/cast_set.h"

#include "common/compiler_util.h" // IWYU pragma: keep

#include "common/config.h"

#include "common/logging.h" // LOG

#include "common/status.h"

#include "core/assert_cast.h"

#include "core/block/block.h"

#include "core/block/column_with_type_and_name.h"

#include "core/column/column_nullable.h"

#include "core/column/column_vector.h"

#include "core/data_type/data_type.h"

#include "core/data_type/data_type_factory.hpp"

#include "core/data_type/data_type_number.h" // IWYU pragma: keep

#include "core/types.h"

#include "exec/common/variant_util.h"

#include "io/fs/file_writer.h"

#include "io/fs/local_file_system.h"

#include "runtime/exec_env.h"

#include "runtime/memory/mem_tracker.h"

#include "service/point_query_executor.h"

#include "storage/data_dir.h"

#include "storage/index/index_file_writer.h"

#include "storage/index/inverted/inverted_index_desc.h"

#include "storage/index/inverted/inverted_index_fs_directory.h"

#include "storage/index/primary_key_index.h"

#include "storage/index/short_key_index.h"

#include "storage/iterator/olap_data_convertor.h"

#include "storage/key_coder.h"

#include "storage/olap_common.h"

#include "storage/partial_update_info.h"

#include "storage/row_cursor.h" // RowCursor // IWYU pragma: keep

#include "storage/rowset/rowset_fwd.h"

#include "storage/rowset/rowset_writer_context.h" // RowsetWriterContext

#include "storage/rowset/segment_creator.h"

#include "storage/segment/column_writer.h" // ColumnWriter

#include "storage/segment/external_col_meta_util.h"

#include "storage/segment/historical_row_retriever.h"

#include "storage/segment/page_io.h"

#include "storage/segment/page_pointer.h"

#include "storage/segment/segment_loader.h"

#include "storage/segment/variant/variant_ext_meta_writer.h"

#include "storage/tablet/base_tablet.h"

#include "storage/tablet/tablet_schema.h"

#include "storage/utils.h"

#include "util/coding.h"

#include "util/debug_points.h"

#include "util/faststring.h"

#include "util/json/path_in_data.h"

#include "util/jsonb/serialize.h"

namespace doris::segment_v2 {

using namespace ErrorCode;

using namespace KeyConsts;

static constexpr const char* k_segment_magic = "D0R1";

static constexpr uint32_t k_segment_magic_length = 4;

inline std::string vertical_segment_writer_mem_tracker_name(uint32_t segment_id) {

    return "VerticalSegmentWriter:Segment-" + std::to_string(segment_id);

}

VerticalSegmentWriter::VerticalSegmentWriter(io::FileWriter* file_writer, uint32_t segment_id,

                                             TabletSchemaSPtr tablet_schema, BaseTabletSPtr tablet,

                                             DataDir* data_dir,

                                             const VerticalSegmentWriterOptions& opts,

                                             IndexFileWriter* index_file_writer)

        : _segment_id(segment_id),

          _tablet_schema(std::move(tablet_schema)),

          _tablet(std::move(tablet)),

          _data_dir(data_dir),

          _opts(opts),

          _file_writer(file_writer),

          _index_file_writer(index_file_writer),

          _mem_tracker(std::make_unique<MemTracker>(

                  vertical_segment_writer_mem_tracker_name(segment_id))),

          _mow_context(std::move(opts.mow_ctx)),

          _block_aggregator(*this) {

    CHECK_NOTNULL(file_writer);

    _num_sort_key_columns = _tablet_schema->num_key_columns();

    _num_short_key_columns = _tablet_schema->num_short_key_columns();

    if (!_is_mow_with_cluster_key()) {

        DCHECK(_num_sort_key_columns >= _num_short_key_columns)

                << ", table_id=" << _tablet_schema->table_id()

                << ", num_key_columns=" << _num_sort_key_columns

                << ", num_short_key_columns=" << _num_short_key_columns

                << ", cluster_key_columns=" << _tablet_schema->cluster_key_uids().size();

    }

    for (size_t cid = 0; cid < _num_sort_key_columns; ++cid) {

        const auto& column = _tablet_schema->column(cid);

        _key_coders.push_back(get_key_coder(column.type()));

        _key_index_size.push_back(cast_set<uint16_t>(column.index_length()));

    }

    // encode the sequence id into the primary key index

    if (_is_mow()) {

        if (_tablet_schema->has_sequence_col()) {

            const auto& column = _tablet_schema->column(_tablet_schema->sequence_col_idx());

            _seq_coder = get_key_coder(column.type());

        }

        // encode the rowid into the primary key index

        if (_is_mow_with_cluster_key()) {

            _rowid_coder = get_key_coder(FieldType::OLAP_FIELD_TYPE_UNSIGNED_INT);

            // primary keys

            _primary_key_coders.swap(_key_coders);

            // cluster keys

            _key_coders.clear();

            _key_index_size.clear();

            _num_sort_key_columns = _tablet_schema->cluster_key_uids().size();

            for (auto cid : _tablet_schema->cluster_key_uids()) {

                const auto& column = _tablet_schema->column_by_uid(cid);

                _key_coders.push_back(get_key_coder(column.type()));

                _key_index_size.push_back(cast_set<uint16_t>(column.index_length()));

            }

        }

    }

}

VerticalSegmentWriter::~VerticalSegmentWriter() {

    _mem_tracker->release(_mem_tracker->consumption());

}

void VerticalSegmentWriter::_init_column_meta(ColumnMetaPB* meta, uint32_t column_id,

                                              const TabletColumn& column) {

    meta->set_column_id(column_id);

    meta->set_type(int(column.type()));

    meta->set_length(cast_set<int32_t>(column.length()));

    meta->set_encoding(DEFAULT_ENCODING);

    meta->set_compression(_opts.compression_type);

    meta->set_is_nullable(column.is_nullable());

    meta->set_default_value(column.default_value());

    meta->set_precision(column.precision());

    meta->set_frac(column.frac());

    if (column.has_path_info()) {

        column.path_info_ptr()->to_protobuf(meta->mutable_column_path_info(),

                                            column.parent_unique_id());

    }

    meta->set_unique_id(column.unique_id());

    for (uint32_t i = 0; i < column.get_subtype_count(); ++i) {

        _init_column_meta(meta->add_children_columns(), column_id, column.get_sub_column(i));

    }

    if (column.is_variant_type()) {

        meta->set_variant_max_subcolumns_count(column.variant_max_subcolumns_count());

        meta->set_variant_enable_doc_mode(column.variant_enable_doc_mode());

    }

    meta->set_result_is_nullable(column.get_result_is_nullable());

    meta->set_function_name(column.get_aggregation_name());

    meta->set_be_exec_version(column.get_be_exec_version());

}

Status VerticalSegmentWriter::_create_column_writer(uint32_t cid, const TabletColumn& column,

                                                    const TabletSchemaSPtr& tablet_schema) {

    ColumnWriterOptions opts;

    opts.meta = _footer.add_columns();

    _init_column_meta(opts.meta, cid, column);

    // now we create zone map for key columns in AGG_KEYS or all column in UNIQUE_KEYS or DUP_KEYS

    // except for columns whose type don't support zone map.

    opts.need_zone_map = column.is_key() || tablet_schema->keys_type() != KeysType::AGG_KEYS;

    opts.need_bloom_filter = column.is_bf_column();

    if (opts.need_bloom_filter) {

        opts.bf_options.fpp =

                tablet_schema->has_bf_fpp() ? tablet_schema->bloom_filter_fpp() : 0.05;

    }

    auto* tablet_index = tablet_schema->get_ngram_bf_index(column.unique_id());

    if (tablet_index) {

        opts.need_bloom_filter = true;

        opts.is_ngram_bf_index = true;

        //narrow convert from int32_t to uint8_t and uint16_t which is dangerous

        auto gram_size = tablet_index->get_gram_size();

        auto gram_bf_size = tablet_index->get_gram_bf_size();

        if (gram_size > 256 || gram_size < 1) {

            return Status::NotSupported("Do not support ngram bloom filter for ngram_size: ",

                                        gram_size);

        }

        if (gram_bf_size > 65535 || gram_bf_size < 64) {

            return Status::NotSupported("Do not support ngram bloom filter for bf_size: ",

                                        gram_bf_size);

        }

        opts.gram_size = cast_set<uint8_t>(gram_size);

        opts.gram_bf_size = cast_set<uint16_t>(gram_bf_size);

    }

    bool skip_inverted_index = false;

    if (_opts.rowset_ctx != nullptr) {

        // skip write inverted index for index compaction column

        skip_inverted_index =

                _opts.rowset_ctx->columns_to_do_index_compaction.contains(column.unique_id());

    }

    // skip write inverted index on load if skip_write_index_on_load is true

    if (_opts.write_type == DataWriteType::TYPE_DIRECT &&

        tablet_schema->skip_write_index_on_load()) {

        skip_inverted_index = true;

    }

    if (!skip_inverted_index) {

        auto inverted_indexs = tablet_schema->inverted_indexs(column);

        if (!inverted_indexs.empty()) {

            opts.inverted_indexes = inverted_indexs;

            opts.need_inverted_index = true;

            DCHECK(_index_file_writer != nullptr);

        }

    }

    opts.index_file_writer = _index_file_writer;

    if (const auto& index = tablet_schema->ann_index(column); index != nullptr) {

        opts.ann_index = index;

        opts.need_ann_index = true;

        DCHECK(_index_file_writer != nullptr);

        opts.index_file_writer = _index_file_writer;

    }

#define DISABLE_INDEX_IF_FIELD_TYPE(TYPE)                     \

    if (column.type() == FieldType::OLAP_FIELD_TYPE_##TYPE) { \

        opts.need_zone_map = false;                           \

        opts.need_bloom_filter = false;                       \

    }

    DISABLE_INDEX_IF_FIELD_TYPE(STRUCT)

    DISABLE_INDEX_IF_FIELD_TYPE(ARRAY)

    DISABLE_INDEX_IF_FIELD_TYPE(JSONB)

    DISABLE_INDEX_IF_FIELD_TYPE(AGG_STATE)

    DISABLE_INDEX_IF_FIELD_TYPE(MAP)

    DISABLE_INDEX_IF_FIELD_TYPE(BITMAP)

    DISABLE_INDEX_IF_FIELD_TYPE(HLL)

    DISABLE_INDEX_IF_FIELD_TYPE(QUANTILE_STATE)

    DISABLE_INDEX_IF_FIELD_TYPE(VARIANT)

#undef DISABLE_INDEX_IF_FIELD_TYPE

#undef CHECK_FIELD_TYPE

    int64_t storage_page_size = _tablet_schema->storage_page_size();

    // storage_page_size must be between 4KB and 10MB.

    if (storage_page_size >= 4096 && storage_page_size <= 10485760) {

        opts.data_page_size = storage_page_size;

    }

    opts.dict_page_size = _tablet_schema->storage_dict_page_size();

    DBUG_EXECUTE_IF("VerticalSegmentWriter._create_column_writer.storage_page_size", {

        auto table_id = DebugPoints::instance()->get_debug_param_or_default<int64_t>(

                "VerticalSegmentWriter._create_column_writer.storage_page_size", "table_id",

                INT_MIN);

        auto target_data_page_size = DebugPoints::instance()->get_debug_param_or_default<int64_t>(

                "VerticalSegmentWriter._create_column_writer.storage_page_size",

                "storage_page_size", INT_MIN);

        if (table_id == INT_MIN || target_data_page_size == INT_MIN) {

            return Status::Error<ErrorCode::INTERNAL_ERROR>(

                    "Debug point parameters missing: either 'table_id' or 'storage_page_size' not "

                    "set.");

        }

        if (table_id == _tablet_schema->table_id() &&

            opts.data_page_size != target_data_page_size) {

            return Status::Error<ErrorCode::INTERNAL_ERROR>(

                    "Mismatch in 'storage_page_size': expected size does not match the current "

                    "data page size. "

                    "Expected: " +

                    std::to_string(target_data_page_size) +

                    ", Actual: " + std::to_string(opts.data_page_size) + ".");

        }

    })

    if (column.is_row_store_column()) {

        // smaller page size for row store column

        auto page_size = _tablet_schema->row_store_page_size();

        opts.data_page_size =

                (page_size > 0) ? page_size : segment_v2::ROW_STORE_PAGE_SIZE_DEFAULT_VALUE;

        // Row store data is already serialized as a single blob. Keep it on plain pages

        // to avoid introducing dictionary pages for the hidden row store column.

        opts.meta->set_encoding(_tablet_schema->binary_plain_encoding_default_impl() ==

                                                BinaryPlainEncodingTypePB::BINARY_PLAIN_ENCODING_V2

                                        ? PLAIN_ENCODING_V2

                                        : PLAIN_ENCODING);

    }

    opts.rowset_ctx = _opts.rowset_ctx;

    opts.file_writer = _file_writer;

    opts.compression_type = _opts.compression_type;

    opts.footer = &_footer;

    opts.input_rs_readers = _opts.rowset_ctx->input_rs_readers;

    opts.encoding_preference = {.integer_type_default_use_plain_encoding =

                                        _tablet_schema->integer_type_default_use_plain_encoding(),

                                .binary_plain_encoding_default_impl =

                                        _tablet_schema->binary_plain_encoding_default_impl()};

    std::unique_ptr<ColumnWriter> writer;

    RETURN_IF_ERROR(ColumnWriter::create(opts, &column, _file_writer, &writer));

    RETURN_IF_ERROR(writer->init());

    _column_writers[cid] = std::move(writer);

    _olap_data_convertor->add_column_data_convertor_at(column, cid);

    return Status::OK();

};

Status VerticalSegmentWriter::init() {

    DCHECK(_column_writers.empty());

    if (_opts.compression_type == UNKNOWN_COMPRESSION) {

        _opts.compression_type = _tablet_schema->compression_type();

    }

    _olap_data_convertor = std::make_unique<OlapBlockDataConvertor>();

    _olap_data_convertor->resize(_tablet_schema->num_columns());

    _column_writers.resize(_tablet_schema->num_columns());

    // we don't need the short key index for unique key merge on write table.

    if (_is_mow()) {

        size_t seq_col_length = 0;

        if (_tablet_schema->has_sequence_col()) {

            seq_col_length =

                    _tablet_schema->column(_tablet_schema->sequence_col_idx()).length() + 1;

        }

        size_t rowid_length = 0;

        if (_is_mow_with_cluster_key()) {

            rowid_length = PrimaryKeyIndexReader::ROW_ID_LENGTH;

            _short_key_index_builder.reset(

                    new ShortKeyIndexBuilder(_segment_id, _opts.num_rows_per_block));

        }

        _primary_key_index_builder.reset(

                new PrimaryKeyIndexBuilder(_file_writer, seq_col_length, rowid_length));

        RETURN_IF_ERROR(_primary_key_index_builder->init());

    } else {

        _short_key_index_builder.reset(

                new ShortKeyIndexBuilder(_segment_id, _opts.num_rows_per_block));

    }

    return Status::OK();

}

void VerticalSegmentWriter::_maybe_invalid_row_cache(const std::string& key) const {

    // Just invalid row cache for simplicity, since the rowset is not visible at present.

    // If we update/insert cache, if load failed rowset will not be visible but cached data

    // will be visible, and lead to inconsistency.

    if (!config::disable_storage_row_cache && _tablet_schema->has_row_store_for_all_columns() &&

        _opts.write_type == DataWriteType::TYPE_DIRECT) {

        // invalidate cache

        RowCache::instance()->erase({_opts.rowset_ctx->tablet_id, key});

    }

}

void VerticalSegmentWriter::_serialize_block_to_row_column(const Block& block) {

    if (block.rows() == 0) {

        return;

    }

    MonotonicStopWatch watch;

    watch.start();

    int row_column_id = 0;

    for (int i = 0; i < _tablet_schema->num_columns(); ++i) {

        if (_tablet_schema->column(i).is_row_store_column()) {

            auto* row_store_column = static_cast<ColumnString*>(

                    block.get_by_position(i).column->assume_mutable_ref().assume_mutable().get());

            row_store_column->clear();

            DataTypeSerDeSPtrs serdes = create_data_type_serdes(block.get_data_types());

            std::unordered_set<int> row_store_cids_set(_tablet_schema->row_columns_uids().begin(),

                                                       _tablet_schema->row_columns_uids().end());

            JsonbSerializeUtil::block_to_jsonb(*_tablet_schema, block, *row_store_column,

                                               cast_set<int>(_tablet_schema->num_columns()), serdes,

                                               row_store_cids_set);

            break;

        }

    }

    VLOG_DEBUG << "serialize , num_rows:" << block.rows() << ", row_column_id:" << row_column_id

               << ", total_byte_size:" << block.allocated_bytes() << ", serialize_cost(us)"

               << watch.elapsed_time() / 1000;

}

Status VerticalSegmentWriter::_probe_key_for_mow(

        std::string key, std::size_t segment_pos, bool have_input_seq_column, bool have_delete_sign,

        const std::vector<RowsetSharedPtr>& specified_rowsets,

        std::vector<std::unique_ptr<SegmentCacheHandle>>& segment_caches,

        bool& has_default_or_nullable, std::vector<bool>& use_default_or_null_flag,

        const std::function<void(const RowLocation& loc)>& found_cb,

        const std::function<Status()>& not_found_cb, PartialUpdateStats& stats) {

    RowLocation loc;

    // save rowset shared ptr so this rowset wouldn't delete

    RowsetSharedPtr rowset;

    auto st = _tablet->lookup_row_key(key, _tablet_schema.get(), have_input_seq_column,

                                      specified_rowsets, &loc, _mow_context->max_version,

                                      segment_caches, &rowset);

    if (st.is<KEY_NOT_FOUND>()) {

        if (!have_delete_sign) {

            RETURN_IF_ERROR(not_found_cb());

        }

        ++stats.num_rows_new_added;

        has_default_or_nullable = true;

        use_default_or_null_flag.emplace_back(true);

        return Status::OK();

    }

    if (!st.ok() && !st.is<KEY_ALREADY_EXISTS>()) {

        LOG(WARNING) << "failed to lookup row key, error: " << st;

        return st;

    }

    // 1. if the delete sign is marked, it means that the value columns of the row will not

    //    be read. So we don't need to read the missing values from the previous rows.

    // 2. the one exception is when there are sequence columns in the table, we need to read

    //    the sequence columns, otherwise it may cause the merge-on-read based compaction

    //    policy to produce incorrect results

    // 3. In flexible partial update, we may delete the existing rows before if there exists

    //    insert after delete in one load. In this case, the insert should also be treated

    //    as newly inserted rows, note that the sequence column value is filled in

    //    BlockAggregator::aggregate_for_insert_after_delete() if this row doesn't specify the sequence column

    if (st.is<KEY_ALREADY_EXISTS>() || (have_delete_sign && !_tablet_schema->has_sequence_col()) ||

        (_opts.rowset_ctx->partial_update_info->is_flexible_partial_update() &&

         _mow_context->delete_bitmap->contains(

                 {loc.rowset_id, loc.segment_id, DeleteBitmap::TEMP_VERSION_COMMON}, loc.row_id))) {

        has_default_or_nullable = true;

        use_default_or_null_flag.emplace_back(true);

    } else {

        // partial update should not contain invisible columns

        use_default_or_null_flag.emplace_back(false);

        _rsid_to_rowset.emplace(rowset->rowset_id(), rowset);

        found_cb(loc);

    }

    if (st.is<KEY_ALREADY_EXISTS>()) {

        // although we need to mark delete current row, we still need to read missing columns

        // for this row, we need to ensure that each column is aligned

        _mow_context->delete_bitmap->add(

                {_opts.rowset_ctx->rowset_id, _segment_id, DeleteBitmap::TEMP_VERSION_COMMON},

                cast_set<uint32_t>(segment_pos));

        ++stats.num_rows_deleted;

    } else {

        _mow_context->delete_bitmap->add(

                {loc.rowset_id, loc.segment_id, DeleteBitmap::TEMP_VERSION_COMMON}, loc.row_id);

        ++stats.num_rows_updated;

    }

    return Status::OK();

}

Status VerticalSegmentWriter::_finalize_column_writer_and_update_meta(size_t cid) {

    RETURN_IF_ERROR(_column_writers[cid]->finish());

    RETURN_IF_ERROR(_column_writers[cid]->write_data());

    auto* column_meta = _column_writers[cid]->get_column_meta();

    column_meta->set_compressed_data_bytes(

            _column_writers[cid]->get_total_compressed_data_pages_bytes());

    column_meta->set_uncompressed_data_bytes(

            _column_writers[cid]->get_total_uncompressed_data_pages_bytes());

    column_meta->set_raw_data_bytes(_column_writers[cid]->get_raw_data_bytes());

    return Status::OK();

}

Status VerticalSegmentWriter::_partial_update_preconditions_check(size_t row_pos,

                                                                  bool is_flexible_update) {

    if (!_is_mow()) {

        auto msg = fmt::format(

                "Can only do partial update on merge-on-write unique table, but found: "

                "keys_type={}, _opts.enable_unique_key_merge_on_write={}, tablet_id={}",

                _tablet_schema->keys_type(), _opts.enable_unique_key_merge_on_write,

                _tablet->tablet_id());

        DCHECK(false) << msg;

        return Status::InternalError<false>(msg);

    }

    if (_opts.rowset_ctx->partial_update_info == nullptr) {

        auto msg =

                fmt::format("partial_update_info should not be nullptr, please check, tablet_id={}",

                            _tablet->tablet_id());

        DCHECK(false) << msg;

        return Status::InternalError<false>(msg);

    }

    if (!is_flexible_update) {

        if (!_opts.rowset_ctx->partial_update_info->is_fixed_partial_update()) {

            auto msg = fmt::format(

                    "in fixed partial update code, but update_mode={}, please check, tablet_id={}",

                    _opts.rowset_ctx->partial_update_info->update_mode(), _tablet->tablet_id());

            DCHECK(false) << msg;

            return Status::InternalError<false>(msg);

        }

    } else {

        if (!_opts.rowset_ctx->partial_update_info->is_flexible_partial_update()) {

            auto msg = fmt::format(

                    "in flexible partial update code, but update_mode={}, please check, "

                    "tablet_id={}",

                    _opts.rowset_ctx->partial_update_info->update_mode(), _tablet->tablet_id());

            DCHECK(false) << msg;

            return Status::InternalError<false>(msg);

        }

    }

    if (row_pos != 0) {

        auto msg = fmt::format("row_pos should be 0, but found {}, tablet_id={}", row_pos,

                               _tablet->tablet_id());

        DCHECK(false) << msg;

        return Status::InternalError<false>(msg);

    }

    return Status::OK();

}

// for partial update, we should do following steps to fill content of block:

// 1. set block data to data convertor, and get all key_column's converted slice

// 2. get pk of input block, and read missing columns

//       2.1 first find key location{rowset_id, segment_id, row_id}

//       2.2 build read plan to read by batch

//       2.3 fill block

// 3. set columns to data convertor and then write all columns

Status VerticalSegmentWriter::_append_block_with_partial_content(RowsInBlock& data,

                                                                 Block& full_block) {

    DBUG_EXECUTE_IF("_append_block_with_partial_content.block", DBUG_BLOCK);

    RETURN_IF_ERROR(_partial_update_preconditions_check(data.row_pos, false));

    // create full block and fill with input columns

    full_block = _tablet_schema->create_block();

    const auto& including_cids = _opts.rowset_ctx->partial_update_info->update_cids;

    size_t input_id = 0;

    for (auto i : including_cids) {

        full_block.replace_by_position(i, data.block->get_by_position(input_id++).column);

    }

    if (_opts.rowset_ctx->write_type != DataWriteType::TYPE_COMPACTION &&

        _tablet_schema->num_variant_columns() > 0) {

        RETURN_IF_ERROR(variant_util::parse_and_materialize_variant_columns(

                full_block, *_tablet_schema, including_cids));

    }

    bool have_input_seq_column = false;

    // write including columns

    std::vector<IOlapColumnDataAccessor*> key_columns;

    IOlapColumnDataAccessor* seq_column = nullptr;

    uint32_t segment_start_pos = 0;

    for (auto cid : including_cids) {

        RETURN_IF_ERROR(_create_column_writer(cid, _tablet_schema->column(cid), _tablet_schema));

        RETURN_IF_ERROR(_olap_data_convertor->set_source_content_with_specifid_columns(

                &full_block, data.row_pos, data.num_rows, std::vector<uint32_t> {cid}));

        // here we get segment column row num before append data.

        segment_start_pos = cast_set<uint32_t>(_column_writers[cid]->get_next_rowid());

        // olap data convertor alway start from id = 0

        auto [status, column] = _olap_data_convertor->convert_column_data(cid);

        if (!status.ok()) {

            return status;

        }

        if (cid < _num_sort_key_columns) {

            key_columns.push_back(column);

        } else if (_tablet_schema->has_sequence_col() &&

                   cid == _tablet_schema->sequence_col_idx()) {

            seq_column = column;

            have_input_seq_column = true;

        }

        RETURN_IF_ERROR(_column_writers[cid]->append(column->get_nullmap(), column->get_data(),

                                                     data.num_rows));

        RETURN_IF_ERROR(_finalize_column_writer_and_update_meta(cid));

        // Don't clear source content for key columns and sequence column here,

        // as they will be used later in _full_encode_keys() and _generate_primary_key_index().

        // They will be cleared at the end of this method.

        bool is_key_column = (cid < _num_sort_key_columns);

        bool is_seq_column = (_tablet_schema->has_sequence_col() &&

                              cid == _tablet_schema->sequence_col_idx() && have_input_seq_column);

        if (!is_key_column && !is_seq_column) {

            _olap_data_convertor->clear_source_content(cid);

        }

    }

    bool has_default_or_nullable = false;

    std::vector<bool> use_default_or_null_flag;

    use_default_or_null_flag.reserve(data.num_rows);

    const auto* delete_signs =

            BaseTablet::get_delete_sign_column_data(full_block, data.row_pos + data.num_rows);

    DBUG_EXECUTE_IF("VerticalSegmentWriter._append_block_with_partial_content.sleep",

                    { sleep(60); })

    const std::vector<RowsetSharedPtr>& specified_rowsets = _mow_context->rowset_ptrs;

    std::vector<std::unique_ptr<SegmentCacheHandle>> segment_caches(specified_rowsets.size());

    FixedReadPlan read_plan;

    // locate rows in base data

    PartialUpdateStats stats;

    for (size_t block_pos = data.row_pos; block_pos < data.row_pos + data.num_rows; block_pos++) {

        // block   segment

        //   2   ->   0

        //   3   ->   1

        //   4   ->   2

        //   5   ->   3

        // here row_pos = 2, num_rows = 4.

        size_t delta_pos = block_pos - data.row_pos;

        size_t segment_pos = segment_start_pos + delta_pos;

        std::string key = _full_encode_keys(key_columns, delta_pos);

        _maybe_invalid_row_cache(key);

        if (have_input_seq_column) {

            _encode_seq_column(seq_column, delta_pos, &key);

        }

        // If the table have sequence column, and the include-cids don't contain the sequence

        // column, we need to update the primary key index builder at the end of this method.

        // At that time, we have a valid sequence column to encode the key with seq col.

        if (!_tablet_schema->has_sequence_col() || have_input_seq_column) {

            RETURN_IF_ERROR(_primary_key_index_builder->add_item(key));

        }

        // mark key with delete sign as deleted.

        bool have_delete_sign = (delete_signs != nullptr && delete_signs[block_pos] != 0);

        auto not_found_cb = [&]() {

            return _opts.rowset_ctx->partial_update_info->handle_new_key(

                    *_tablet_schema, [&]() -> std::string {

                        return data.block->dump_one_line(block_pos,

                                                         cast_set<int>(_num_sort_key_columns));

                    });

        };

        auto update_read_plan = [&](const RowLocation& loc) {

            read_plan.prepare_to_read(loc, segment_pos);

        };

        RETURN_IF_ERROR(_probe_key_for_mow(std::move(key), segment_pos, have_input_seq_column,

                                           have_delete_sign, specified_rowsets, segment_caches,

                                           has_default_or_nullable, use_default_or_null_flag,

                                           update_read_plan, not_found_cb, stats));

    }

    CHECK_EQ(use_default_or_null_flag.size(), data.num_rows);

    if (config::enable_merge_on_write_correctness_check) {

        _tablet->add_sentinel_mark_to_delete_bitmap(_mow_context->delete_bitmap.get(),

                                                    *_mow_context->rowset_ids);

    }

    // read to fill full_block

    RETURN_IF_ERROR(read_plan.fill_missing_columns(

            _opts.rowset_ctx->make_historical_row_retriever_context(), _rsid_to_rowset,

            *_tablet_schema, full_block, use_default_or_null_flag, has_default_or_nullable,

            segment_start_pos, data.block));

    if (_tablet_schema->num_variant_columns() > 0) {

        RETURN_IF_ERROR(variant_util::parse_and_materialize_variant_columns(

                full_block, *_tablet_schema, _opts.rowset_ctx->partial_update_info->missing_cids));

    }

    // row column should be filled here

    // convert block to row store format

    _serialize_block_to_row_column(full_block);

    // convert missing columns and send to column writer

    const auto& missing_cids = _opts.rowset_ctx->partial_update_info->missing_cids;

    for (auto cid : missing_cids) {

        RETURN_IF_ERROR(_create_column_writer(cid, _tablet_schema->column(cid), _tablet_schema));

        RETURN_IF_ERROR(_olap_data_convertor->set_source_content_with_specifid_columns(

                &full_block, data.row_pos, data.num_rows, std::vector<uint32_t> {cid}));

        auto [status, column] = _olap_data_convertor->convert_column_data(cid);

        if (!status.ok()) {

            return status;

        }

        if (_tablet_schema->has_sequence_col() && !have_input_seq_column &&

            cid == _tablet_schema->sequence_col_idx()) {

            DCHECK_EQ(seq_column, nullptr);

            seq_column = column;

        }

        RETURN_IF_ERROR(_column_writers[cid]->append(column->get_nullmap(), column->get_data(),

                                                     data.num_rows));

        RETURN_IF_ERROR(_finalize_column_writer_and_update_meta(cid));

        // Don't clear source content for sequence column here if it will be used later

        // in _generate_primary_key_index(). It will be cleared at the end of this method.

        bool is_seq_column = (_tablet_schema->has_sequence_col() && !have_input_seq_column &&

                              cid == _tablet_schema->sequence_col_idx());

        if (!is_seq_column) {

            _olap_data_convertor->clear_source_content(cid);

        }

    }

    _num_rows_updated += stats.num_rows_updated;

    _num_rows_deleted += stats.num_rows_deleted;

    _num_rows_new_added += stats.num_rows_new_added;

    _num_rows_filtered += stats.num_rows_filtered;

    if (_tablet_schema->has_sequence_col() && !have_input_seq_column) {

        DCHECK_NE(seq_column, nullptr);

        if (_num_rows_written != data.row_pos ||

            _primary_key_index_builder->num_rows() != _num_rows_written) {

            return Status::InternalError(

                    "Correctness check failed, _num_rows_written: {}, row_pos: {}, primary key "

                    "index builder num rows: {}",

                    _num_rows_written, data.row_pos, _primary_key_index_builder->num_rows());

        }

        RETURN_IF_ERROR(_generate_primary_key_index(_key_coders, key_columns, seq_column,

                                                    data.num_rows, false));

    }

    _num_rows_written += data.num_rows;

    DCHECK_EQ(_primary_key_index_builder->num_rows(), _num_rows_written)

            << "primary key index builder num rows(" << _primary_key_index_builder->num_rows()

            << ") not equal to segment writer's num rows written(" << _num_rows_written << ")";

    _olap_data_convertor->clear_source_content();

    return Status::OK();

}

Status VerticalSegmentWriter::_append_block_with_flexible_partial_content(RowsInBlock& data,

                                                                          Block& full_block) {

    RETURN_IF_ERROR(_partial_update_preconditions_check(data.row_pos, true));

    // data.block has the same schema with full_block

    DCHECK(data.block->columns() == _tablet_schema->num_columns());

    // create full block and fill with sort key columns

    full_block = _tablet_schema->create_block();

    // Use _num_rows_written instead of creating column writer 0, since all column writers

    // should have the same row count, which equals _num_rows_written.

    uint32_t segment_start_pos = cast_set<uint32_t>(_num_rows_written);

    DCHECK(_tablet_schema->has_skip_bitmap_col());

    auto skip_bitmap_col_idx = _tablet_schema->skip_bitmap_col_idx();

    bool has_default_or_nullable = false;

    std::vector<bool> use_default_or_null_flag;

    use_default_or_null_flag.reserve(data.num_rows);

    int32_t seq_map_col_unique_id = _opts.rowset_ctx->partial_update_info->sequence_map_col_uid();

    bool schema_has_sequence_col = _tablet_schema->has_sequence_col();

    DBUG_EXECUTE_IF("VerticalSegmentWriter._append_block_with_flexible_partial_content.sleep",

                    { sleep(60); })

    const std::vector<RowsetSharedPtr>& specified_rowsets = _mow_context->rowset_ptrs;

    std::vector<std::unique_ptr<SegmentCacheHandle>> segment_caches(specified_rowsets.size());

    // Ensure all primary key column writers and sequence column writer are created before

    // aggregate_for_flexible_partial_update, because it internally calls convert_pk_columns

    // and convert_seq_column which need the convertors in _olap_data_convertor

    for (uint32_t cid = 0; cid < _tablet_schema->num_key_columns(); ++cid) {

        RETURN_IF_ERROR(_create_column_writer(cid, _tablet_schema->column(cid), _tablet_schema));

    }

    if (schema_has_sequence_col) {

        uint32_t cid = _tablet_schema->sequence_col_idx();

        RETURN_IF_ERROR(_create_column_writer(cid, _tablet_schema->column(cid), _tablet_schema));

    }

    // 1. aggregate duplicate rows in block

    RETURN_IF_ERROR(_block_aggregator.aggregate_for_flexible_partial_update(

            const_cast<Block*>(data.block), data.num_rows, specified_rowsets, segment_caches));

    if (data.block->rows() != data.num_rows) {

        data.num_rows = data.block->rows();

        _olap_data_convertor->clear_source_content();

    }

    // 2. encode primary key columns

    // we can only encode primary key columns currently becasue all non-primary columns in flexible partial update

    // can have missing cells

    std::vector<IOlapColumnDataAccessor*> key_columns {};

    RETURN_IF_ERROR(_block_aggregator.convert_pk_columns(const_cast<Block*>(data.block),

                                                         data.row_pos, data.num_rows, key_columns));

    // 3. encode sequence column

    // We encode the seguence column even thought it may have invalid values in some rows because we need to

    // encode the value of sequence column in key for rows that have a valid value in sequence column during

    // lookup_raw_key. We will encode the sequence column again at the end of this method. At that time, we have

    // a valid sequence column to encode the key with seq col.

    IOlapColumnDataAccessor* seq_column {nullptr};

    RETURN_IF_ERROR(_block_aggregator.convert_seq_column(const_cast<Block*>(data.block),

                                                         data.row_pos, data.num_rows, seq_column));

    std::vector<BitmapValue>* skip_bitmaps = &(

            assert_cast<ColumnBitmap*>(

                    data.block->get_by_position(skip_bitmap_col_idx).column->assume_mutable().get())

                    ->get_data());

    const auto* delete_signs =

            BaseTablet::get_delete_sign_column_data(*data.block, data.row_pos + data.num_rows);

    DCHECK(delete_signs != nullptr);

    for (std::size_t cid {0}; cid < _tablet_schema->num_key_columns(); cid++) {

        full_block.replace_by_position(cid, data.block->get_by_position(cid).column);

    }

    // 4. write primary key columns data

    for (std::size_t cid {0}; cid < _tablet_schema->num_key_columns(); cid++) {

        const auto& column = key_columns[cid];

        DCHECK(_column_writers[cid]->get_next_rowid() == _num_rows_written);

        RETURN_IF_ERROR(_column_writers[cid]->append(column->get_nullmap(), column->get_data(),

                                                     data.num_rows));

        DCHECK(_column_writers[cid]->get_next_rowid() == _num_rows_written + data.num_rows);

        RETURN_IF_ERROR(_finalize_column_writer_and_update_meta(cid));

    }

    // 5. genreate read plan

    FlexibleReadPlan read_plan {_tablet_schema->has_row_store_for_all_columns()};

    PartialUpdateStats stats;

    RETURN_IF_ERROR(_generate_flexible_read_plan(

            read_plan, data, segment_start_pos, schema_has_sequence_col, seq_map_col_unique_id,

            skip_bitmaps, key_columns, seq_column, delete_signs, specified_rowsets, segment_caches,

            has_default_or_nullable, use_default_or_null_flag, stats));

    CHECK_EQ(use_default_or_null_flag.size(), data.num_rows);

    if (config::enable_merge_on_write_correctness_check) {

        _tablet->add_sentinel_mark_to_delete_bitmap(_mow_context->delete_bitmap.get(),

                                                    *_mow_context->rowset_ids);

    }

    // 6. read according plan to fill full_block

    RETURN_IF_ERROR(read_plan.fill_non_primary_key_columns(

            _opts.rowset_ctx->make_historical_row_retriever_context(), _rsid_to_rowset,

            *_tablet_schema, full_block, use_default_or_null_flag, has_default_or_nullable,

            segment_start_pos, cast_set<uint32_t>(data.row_pos), data.block, skip_bitmaps));

    // TODO(bobhan1): should we replace the skip bitmap column with empty bitmaps to reduce storage occupation?

    // this column is not needed in read path for merge-on-write table

    // 7. fill row store column

    _serialize_block_to_row_column(full_block);

    std::vector<uint32_t> column_ids;

    for (uint32_t i = 0; i < _tablet_schema->num_columns(); ++i) {

        column_ids.emplace_back(i);

    }

    if (_opts.rowset_ctx->write_type != DataWriteType::TYPE_COMPACTION &&

        _tablet_schema->num_variant_columns() > 0) {

        RETURN_IF_ERROR(variant_util::parse_and_materialize_variant_columns(

                full_block, *_tablet_schema, column_ids));

    }

    // 8. encode and write all non-primary key columns(including sequence column if exists)

    for (auto cid = _tablet_schema->num_key_columns(); cid < _tablet_schema->num_columns(); cid++) {

        if (cid != _tablet_schema->sequence_col_idx()) {

            RETURN_IF_ERROR(_create_column_writer(cast_set<uint32_t>(cid),

                                                  _tablet_schema->column(cid), _tablet_schema));

        }

        RETURN_IF_ERROR(_olap_data_convertor->set_source_content_with_specifid_column(

                full_block.get_by_position(cid), data.row_pos, data.num_rows,

                cast_set<uint32_t>(cid)));

        auto [status, column] = _olap_data_convertor->convert_column_data(cid);

        if (!status.ok()) {

            return status;

        }

        if (cid == _tablet_schema->sequence_col_idx()) {

            // should use the latest encoded sequence column to build the primary index

            seq_column = column;

        }

        DCHECK(_column_writers[cid]->get_next_rowid() == _num_rows_written);

        RETURN_IF_ERROR(_column_writers[cid]->append(column->get_nullmap(), column->get_data(),

                                                     data.num_rows));

        DCHECK(_column_writers[cid]->get_next_rowid() == _num_rows_written + data.num_rows);

        RETURN_IF_ERROR(_finalize_column_writer_and_update_meta(cid));

    }

    _num_rows_updated += stats.num_rows_updated;

    _num_rows_deleted += stats.num_rows_deleted;

    _num_rows_new_added += stats.num_rows_new_added;

    _num_rows_filtered += stats.num_rows_filtered;

    if (_num_rows_written != data.row_pos ||

        _primary_key_index_builder->num_rows() != _num_rows_written) {

        return Status::InternalError(

                "Correctness check failed, _num_rows_written: {}, row_pos: {}, primary key "

                "index builder num rows: {}",

                _num_rows_written, data.row_pos, _primary_key_index_builder->num_rows());

    }

    // 9. build primary key index

    RETURN_IF_ERROR(_generate_primary_key_index(_key_coders, key_columns, seq_column, data.num_rows,

                                                false));

    _num_rows_written += data.num_rows;

    DCHECK_EQ(_primary_key_index_builder->num_rows(), _num_rows_written)

            << "primary key index builder num rows(" << _primary_key_index_builder->num_rows()

            << ") not equal to segment writer's num rows written(" << _num_rows_written << ")";

    _olap_data_convertor->clear_source_content();

    return Status::OK();

}

Status VerticalSegmentWriter::_generate_encoded_default_seq_value(const TabletSchema& tablet_schema,

                                                                  const PartialUpdateInfo& info,

                                                                  std::string* encoded_value) {

    const auto& seq_column = tablet_schema.column(tablet_schema.sequence_col_idx());

    auto block = tablet_schema.create_block_by_cids(

            {cast_set<uint32_t>(tablet_schema.sequence_col_idx())});

    if (seq_column.has_default_value()) {

        auto idx = tablet_schema.sequence_col_idx() - tablet_schema.num_key_columns();

        const auto& default_value = info.default_values[idx];

        StringRef str {default_value};

        RETURN_IF_ERROR(block.get_by_position(0).type->get_serde()->default_from_string(

                str, *block.get_by_position(0).column->assume_mutable().get()));

    } else {

        block.get_by_position(0).column->assume_mutable()->insert_default();

    }

    DCHECK_EQ(block.rows(), 1);

    auto olap_data_convertor = std::make_unique<OlapBlockDataConvertor>();

    olap_data_convertor->add_column_data_convertor(seq_column);

    olap_data_convertor->set_source_content(&block, 0, 1);

    auto [status, column] = olap_data_convertor->convert_column_data(0);

    if (!status.ok()) {

        return status;

    }

    // include marker

    _encode_seq_column(column, 0, encoded_value);

    return Status::OK();

}

Status VerticalSegmentWriter::_generate_flexible_read_plan(

        FlexibleReadPlan& read_plan, RowsInBlock& data, size_t segment_start_pos,

        bool schema_has_sequence_col, int32_t seq_map_col_unique_id,

        std::vector<BitmapValue>* skip_bitmaps,

        const std::vector<IOlapColumnDataAccessor*>& key_columns,

        IOlapColumnDataAccessor* seq_column, const signed char* delete_signs,

        const std::vector<RowsetSharedPtr>& specified_rowsets,

        std::vector<std::unique_ptr<SegmentCacheHandle>>& segment_caches,

        bool& has_default_or_nullable, std::vector<bool>& use_default_or_null_flag,

        PartialUpdateStats& stats) {

    int32_t delete_sign_col_unique_id =

            _tablet_schema->column(_tablet_schema->delete_sign_idx()).unique_id();

    int32_t seq_col_unique_id =

            (_tablet_schema->has_sequence_col()

                     ? _tablet_schema->column(_tablet_schema->sequence_col_idx()).unique_id()

                     : -1);

    for (size_t block_pos = data.row_pos; block_pos < data.row_pos + data.num_rows; block_pos++) {