// 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 <algorithm>

#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_string.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 "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/mow/historical_row_fetcher.h"

#include "storage/mow/key_probe.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/encoding_info.h"

#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/transform/block_transform.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;

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

}

static ColumnBitmap* get_mutable_skip_bitmap_column(Block* block, size_t skip_bitmap_col_idx) {

    auto skip_bitmap_column =

            IColumn::mutate(std::move(block->get_by_position(skip_bitmap_col_idx).column));

    auto* skip_bitmap_column_ptr = assert_cast<ColumnBitmap*>(skip_bitmap_column.get());

    block->replace_by_position(skip_bitmap_col_idx, std::move(skip_bitmap_column));

    return skip_bitmap_column_ptr;

}

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

          _key_encoder(*_tablet_schema, _is_mow()),

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

          _block_aggregator(*this) {

    CHECK_NOTNULL(file_writer);

    _num_short_key_columns = _tablet_schema->num_short_key_columns();

}

VerticalSegmentWriter::~VerticalSegmentWriter() {

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

}

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

                                              const TabletColumn& column,

                                              const ColumnWriterOptions& opts) {

    meta->set_column_id(column_id);

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

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

    meta->set_encoding(EncodingInfo::resolve_default_encoding(opts.storage_format, column));

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

    }

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

    opts.storage_format = tablet_schema->storage_format();

    _init_column_meta(opts.meta, cid, column, opts);

    // 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; encoding is already set to PLAIN /

        // PLAIN_V2 by _init_column_meta via resolve_default_encoding().

        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;

    }

    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;

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

}

Status VerticalSegmentWriter::_append_row_store_column(const Block& block, size_t row_pos,

                                                       size_t num_rows, uint32_t cid) {

    DCHECK(_tablet_schema->column(cid).is_row_store_column());

    if (num_rows == 0) {

        return Status::OK();

    }

    DCHECK_LE(row_pos + num_rows, block.rows());

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

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

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

    size_t end_pos = row_pos + num_rows;

    size_t batch_rows = _opts.num_rows_per_block;

    static constexpr size_t kRowStoreBatchBytes = 4 * 1024 * 1024;

    DCHECK_GT(batch_rows, 0);

    for (size_t pos = row_pos; pos < end_pos;) {

        size_t max_rows = std::min(batch_rows, end_pos - pos);

        auto row_column = ColumnString::create();

        auto* row_store_column = row_column.get();

        size_t rows = JsonbSerializeUtil::block_to_jsonb(

                *_tablet_schema, block, *row_store_column,

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

                max_rows, kRowStoreBatchBytes);

        DCHECK_GT(rows, 0);

        auto typed_column = block.get_by_position(cid);

        typed_column.column = std::move(row_column);

        RETURN_IF_ERROR(_olap_data_convertor->set_source_content_with_specifid_column(

                typed_column, 0, rows, cid));

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

        RETURN_IF_ERROR(status);

        RETURN_IF_ERROR(

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

        _olap_data_convertor->clear_source_content(cid);

        pos += rows;

    }

    return Status::OK();

}

Status VerticalSegmentWriter::_append_generated_column(const DerivedColumnGenerator& generator,

                                                       const Block& block, size_t row_pos,

                                                       size_t num_rows, uint32_t cid) {

    if (num_rows == 0) {

        return Status::OK();

    }

    DCHECK_LE(row_pos + num_rows, block.rows());

    size_t end_pos = row_pos + num_rows;

    size_t batch_rows = _opts.num_rows_per_block;

    static constexpr size_t kDerivedColumnBatchBytes = 4 * 1024 * 1024;

    DCHECK_GT(batch_rows, 0);

    for (size_t pos = row_pos; pos < end_pos;) {

        size_t max_rows = std::min(batch_rows, end_pos - pos);

        auto generated_column = block.get_by_position(cid).column->clone_empty();

        size_t rows = generator.generate(block, pos, max_rows, kDerivedColumnBatchBytes,

                                         generated_column.get());

        DCHECK_GT(rows, 0);

        auto typed_column = block.get_by_position(cid);

        typed_column.column = std::move(generated_column);

        RETURN_IF_ERROR(_olap_data_convertor->set_source_content_with_specifid_column(

                typed_column, 0, rows, cid));

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

        RETURN_IF_ERROR(status);

        RETURN_IF_ERROR(

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

        _olap_data_convertor->clear_source_content(cid);

        pos += rows;

    }

    return Status::OK();

}

Status VerticalSegmentWriter::_probe_key_for_mow(

        const MowKeyProbe& probe, 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, const RowsetSharedPtr& rowset)>& found_cb,

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

    ProbeOutcome outcome =

            DORIS_TRY(probe.probe(key, segment_pos, have_input_seq_column, have_delete_sign,

                                  specified_rowsets, segment_caches, stats));

    if (outcome.result == KeyProbeResult::NOT_FOUND) {

        if (!have_delete_sign) {

            RETURN_IF_ERROR(not_found_cb());

        }

        has_default_or_nullable = true;

        use_default_or_null_flag.emplace_back(true);

        return Status::OK();

    }

    if (outcome.use_default_or_null) {

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

        found_cb(outcome.loc, outcome.rowset);

    }

    return Status::OK();

}

Status VerticalSegmentWriter::_check_column_writer_disk_capacity(size_t cid) {

    if (_data_dir != nullptr &&

        _data_dir->reach_capacity_limit(_column_writers[cid]->estimate_buffer_size())) {

        return Status::Error<DISK_REACH_CAPACITY_LIMIT>("disk {} exceed capacity limit.",

                                                        _data_dir->path_hash());

    }

    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 < _key_encoder.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 for key encoding and _generate_primary_key_index().

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

        bool is_key_column = (cid < _key_encoder.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());

    MowKeyProbe probe = MowKeyProbe::for_partial_update(

            _tablet.get(), _tablet_schema.get(), _tablet_schema->has_sequence_col(), _mow_context,

            _opts.rowset_ctx->rowset_id, _segment_id, /*flexible=*/false);

    // owns the rowset pins and the read plan for the historical read below

    HistoricalRowFetcher fetcher {_opts.rowset_ctx->make_historical_row_retriever_context()};

    // 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 = encode_mow_key_invalidate_cache(

                _key_encoder, key_columns, seq_column, delta_pos, have_input_seq_column,

                _opts.rowset_ctx->tablet_id, *_tablet_schema, _opts.write_type);

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

                    });

        };

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

            // keep the rowset alive until the historical read below is done

            fetcher.pin_rowset(rowset);

            fetcher.plan_fixed_read(loc, segment_pos);

        };

        RETURN_IF_ERROR(_probe_key_for_mow(

                probe, 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(fetcher.fill_missing_columns(*_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));

    }

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

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

            RETURN_IF_ERROR(_append_row_store_column(full_block, data.row_pos, data.num_rows, cid));

            RETURN_IF_ERROR(_finalize_column_writer_and_update_meta(cid));

            continue;

        }

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

    auto* mutable_block = const_cast<Block*>(data.block);

    std::vector<BitmapValue>* skip_bitmaps =

            &get_mutable_skip_bitmap_column(mutable_block, skip_bitmap_col_idx)->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

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

        if (!_tablet_schema->column(cid).is_row_store_column()) {

            continue;

        }

        RETURN_IF_ERROR(_create_column_writer(cast_set<uint32_t>(cid), _tablet_schema->column(cid),

                                              _tablet_schema));

        RETURN_IF_ERROR(_append_row_store_column(full_block, data.row_pos, data.num_rows,

                                                 cast_set<uint32_t>(cid)));

        RETURN_IF_ERROR(_finalize_column_writer_and_update_meta(cid));

    }

    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 (_tablet_schema->column(cid).is_row_store_column()) {

            continue;

        }

        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_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->assert_mutable().get()));

    } else {

        block.get_by_position(0).column->assert_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

    _key_encoder.append_seq_suffix(encoded_value, column, 0);

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

    MowKeyProbe probe = MowKeyProbe::for_partial_update(

            _tablet.get(), _tablet_schema.get(), _tablet_schema->has_sequence_col(), _mow_context,

            _opts.rowset_ctx->rowset_id, _segment_id, /*flexible=*/true);

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

        size_t delta_pos = block_pos - data.row_pos;

        size_t segment_pos = segment_start_pos + delta_pos;

        auto& skip_bitmap = skip_bitmaps->at(block_pos);

        bool row_has_sequence_col =

                (schema_has_sequence_col && !skip_bitmap.contains(seq_col_unique_id));

        std::string key = encode_mow_key_invalidate_cache(

                _key_encoder, key_columns, seq_column, delta_pos, row_has_sequence_col,

                _opts.rowset_ctx->tablet_id, *_tablet_schema, _opts.write_type);

        // mark key with delete sign as deleted.

        bool have_delete_sign =

                (!skip_bitmap.contains(delete_sign_col_unique_id) && 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>(_key_encoder.num_sort_key_columns()));

                    },

                    &skip_bitmap);

        };

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

            // the flexible fill still reads through the writer's pin map, which the block

            // aggregator also feeds

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

            read_plan.prepare_to_read(loc, segment_pos, skip_bitmap);

        };

        RETURN_IF_ERROR(_probe_key_for_mow(probe, std::move(key), segment_pos, row_has_sequence_col,

                                           have_delete_sign, specified_rowsets, segment_caches,

                                           has_default_or_nullable, use_default_or_null_flag,

                                           update_read_plan, not_found_cb, stats));

    }

    return Status::OK();

}

Status VerticalSegmentWriter::batch_block(const Block* block, size_t row_pos, size_t num_rows) {

    if (_opts.rowset_ctx->partial_update_info &&

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

        _opts.write_type == DataWriteType::TYPE_DIRECT &&

        !_opts.rowset_ctx->is_transient_rowset_writer) {

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

            if (block->columns() != _tablet_schema->num_columns()) {

                return Status::InvalidArgument(

                        "illegal flexible partial update block columns, block columns = {}, "

                        "tablet_schema columns = {}",

                        block->dump_structure(), _tablet_schema->dump_structure());

            }

        } else {

            if (block->columns() < _tablet_schema->num_key_columns() ||

                block->columns() >= _tablet_schema->num_columns()) {

                return Status::InvalidArgument(fmt::format(

                        "illegal partial update block columns: {}, num key columns: {}, total "

                        "schema columns: {}",

                        block->columns(), _tablet_schema->num_key_columns(),

                        _tablet_schema->num_columns()));

            }

        }

    } else if (block->columns() != _tablet_schema->num_columns()) {

        return Status::InvalidArgument(

                "illegal block columns, block columns = {}, tablet_schema columns = {}",

                block->dump_structure(), _tablet_schema->dump_structure());

    }

    _batched_blocks.emplace_back(block, row_pos, num_rows);

    return Status::OK();

}

Status VerticalSegmentWriter::write_batch() {

    if (_opts.rowset_ctx->partial_update_info &&

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

        _opts.write_type == DataWriteType::TYPE_DIRECT &&

        !_opts.rowset_ctx->is_transient_rowset_writer) {

        bool is_flexible_partial_update =

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

        Block full_block;

        for (auto& data : _batched_blocks) {

            if (is_flexible_partial_update) {

                RETURN_IF_ERROR(_append_block_with_flexible_partial_content(data, full_block));

            } else {

                RETURN_IF_ERROR(_append_block_with_partial_content(data, full_block));

            }

        }

        return Status::OK();

    }

    // The transform chain already validated, parsed variants and decided the derived

    // (row-store) column; this writer only pumps the generator in bounded batches.

    if (_derived_column.second) {

        const auto& [cid, generator] = _derived_column;

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

        for (auto& data : _batched_blocks) {

            RETURN_IF_ERROR(_append_generated_column(*generator, *data.block, data.row_pos,

                                                     data.num_rows, cid));