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

#pragma once

#include <gen_cpp/olap_file.pb.h>

#include <gen_cpp/segment_v2.pb.h>

#include <cstddef>

#include <cstdint>

#include <map>

#include <memory> // unique_ptr

#include <string>

#include <unordered_set>

#include <utility>

#include <vector>

#include "common/status.h" // Status

#include "storage/index/index_file_writer.h"

#include "storage/key/row_key_encoder.h"

#include "storage/olap_define.h"

#include "storage/segment/column_writer.h"

#include "storage/segment/segment_index_file_cache_loader.h"

#include "storage/tablet/tablet.h"

#include "storage/tablet/tablet_schema.h"

#include "util/faststring.h"

#include "util/slice.h"

namespace doris {

class Block;

// TODO(lingbin): Should be a conf that can be dynamically adjusted, or a member in the context

const uint32_t MAX_SEGMENT_SIZE = static_cast<uint32_t>(OLAP_MAX_COLUMN_SEGMENT_FILE_SIZE *

                                                        OLAP_COLUMN_FILE_SEGMENT_SIZE_SCALE);

class DataDir;

class MemTracker;

class ShortKeyIndexBuilder;

class PrimaryKeyIndexBuilder;

struct RowsetWriterContext;

namespace io {

class FileWriter;

class FileSystem;

} // namespace io

namespace segment_v2 {

class IndexFileWriter;

class VariantStatsCaculator;

class DerivedColumnGenerator;

// Matches block_transform.h: at most one derived column (the row-store column)

// for each flush, held as a {cid, generator} pair; null generator means none.

using DerivedColumn = std::pair<uint32_t, std::shared_ptr<const DerivedColumnGenerator>>;

struct VerticalSegmentWriterOptions {

    uint32_t num_rows_per_block = 1024;

    // Caps the rows in one segment. Only unit tests set it, to control how many

    // segments a write produces; production leaves it unlimited.

    uint32_t max_rows_per_segment = UINT32_MAX;

    bool enable_unique_key_merge_on_write = false;

    CompressionTypePB compression_type = UNKNOWN_COMPRESSION;

    RowsetWriterContext* rowset_ctx = nullptr;

    DataWriteType write_type = DataWriteType::TYPE_DEFAULT;

};

// A segment is written one column group at a time:

//

//      for (column_group : column_groups) {

//          writer.init(column_group, has_key);

//          writer.append_block(block, ...);   // any number of times

//          writer.finalize_columns(&index_size);

//      }

//      writer.finalize_footer(&file_size);

//

// Vertical compaction and segcompaction write several groups; init() is

// shorthand for one whole-schema group. finalize_columns() writes the

// data pages the group still buffers, then its indexes, and closes the group.

//

// write_block() takes a whole-schema group's rows all at once and writes each

// column's pages as soon as that column is done, so it opens the group itself

// and needs no init():

//

//      writer.set_derived_column(derived_column);   // only if there is one

//      writer.write_block(block, 0, num_rows);      // opens the group itself

//      writer.finalize_columns(&index_size);

//      writer.finalize_footer(&file_size);

class VerticalSegmentWriter {

public:

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

                                   TabletSchemaSPtr tablet_schema, BaseTabletSPtr tablet,

                                   DataDir* data_dir, const VerticalSegmentWriterOptions& opts,

                                   IndexFileWriter* index_file_writer);

    ~VerticalSegmentWriter();

    VerticalSegmentWriter(const VerticalSegmentWriter&) = delete;

    const VerticalSegmentWriter& operator=(const VerticalSegmentWriter&) = delete;

    Status init();

    // Opens a column group; the footer keeps every group's column metas. Pass

    // has_key=true for the group with the key columns, which comes first. A

    // keyless schema has no such group, so its first value group takes it.

    Status init(const std::vector<uint32_t>& col_ids, bool has_key);

    // Feeds the open group. Its column writers buffer the data pages until

    // finalize_columns(), so more rows can always follow.

    Status append_block(const Block* block, size_t row_pos, size_t num_rows);

    // Opens a whole-schema group and writes it in one call, one column at a time,

    // so only one column's pages are in memory.

    Status write_block(const Block* block, size_t row_pos, size_t num_rows);

    // Sets the row-store column for the next write_block, which pumps it from its

    // generator in batches instead of reading it from the block. append_block

    // ignores it.

    void set_derived_column(DerivedColumn derived_column) {

        _derived_column = std::move(derived_column);

    }

    [[nodiscard]] std::string data_dir_path() const {

        return _data_dir == nullptr ? "" : _data_dir->path();

    }

    // rows fed to the open group; zero again once the group is closed

    [[nodiscard]] uint32_t num_rows_written() const { return _num_rows_written; }

    // the segment's row count, settled when the group holding the key columns is finalized

    [[nodiscard]] uint32_t row_count() const { return _row_count; }

    [[nodiscard]] uint32_t segment_id() const { return _segment_id; }

    // Size readings the append_block feed rolls a segment on: how many more

    // rows fit, what the segment holds right now, and the cluster key MOW

    // primary keys still waiting for the finalize-time sort.

    int64_t max_row_to_add(size_t row_avg_size_in_bytes);

    uint64_t estimate_segment_size();

    [[nodiscard]] uint64_t primary_keys_size() const { return _primary_keys_size; }

    // Closes the open group: its buffered data pages (skipped after write_block),

    // its index sections, and the key indexes if it holds the key columns. Ends

    // with clear(), so the next init() can follow.

    Status finalize_columns(uint64_t* index_size);

    // Ends the segment, once every group is in.

    Status finalize_footer(uint64_t* segment_file_size,

                           SegmentIndexFileCacheInfo* index_file_cache_info = nullptr);

    Slice min_encoded_key();

    Slice max_encoded_key();

    void clear();

    Status close_inverted_index(int64_t* inverted_index_file_size) {

        // no inverted index

        if (_index_file_writer == nullptr) {

            *inverted_index_file_size = 0;

            return Status::OK();

        }

        RETURN_IF_ERROR(_index_file_writer->begin_close());

        *inverted_index_file_size = _index_file_writer->get_index_file_total_size();

        return Status::OK();

    }

private:

    friend class TestVerticalSegmentWriter;

    void _abandon_index_staging();

    void _init_column_meta(ColumnMetaPB* meta, uint32_t column_id, const TabletColumn& column,

                           const ColumnWriterOptions& opts);

    // pos is the column's place in the open group, which is what the column

    // writers are indexed by; cid is its schema column id, which only the

    // footer meta records.

    Status _create_column_writer(size_t pos, uint32_t cid, const TabletSchemaSPtr& tablet_schema);

    // Opens a column group without its column writers: init() then creates

    // them all at once, write_block() one at a time.

    Status _open_group(const std::vector<uint32_t>& col_ids, bool has_key);

    Status _create_writers(const TabletSchemaSPtr& tablet_schema,

                           const std::vector<uint32_t>& col_ids);

    std::vector<uint32_t> _all_column_ids() const;

    // bytes the key index builders hold

    uint64_t _key_index_size();

    // cluster key MOW: sorts the collected primary keys, hands them to the index

    // builder and releases them. The builder decides what reaches the file.

    Status _sort_primary_keys_into_index();

    Status _set_row_count();

    // Closes the open group's column writers and writes their data pages, after

    // checking the data dir has room for them.

    Status _finalize_columns_data();

    Status _write_data();

    Status _write_ordinal_index();

    Status _write_zone_map();

    Status _write_inverted_index();

    Status _write_ann_index();

    Status _write_bloom_filter_index();

    Status _write_short_key_index();

    Status _write_primary_key_index();

    Status _write_footer();

    Status _write_raw_data(const std::vector<Slice>& slices);

    void _set_min_max_key(const Slice& key);

    void _set_min_key(const Slice& key);

    void _set_max_key(const Slice& key);

    Status _append_generated_column(const DerivedColumnGenerator& generator, const Block& block,

                                    size_t row_pos, size_t num_rows, uint32_t cid);

    // Remembers the column if it is a cluster key. _generate_key_index

    // needs them in cluster key order, which is not the group's order.

    void _collect_cluster_key_column(uint32_t cid, const IColumn* column,

                                     std::map<uint32_t, const IColumn*>* cluster_key_columns);

    // The key columns are whole block columns, so rows are numbered from the

    // block's start: the batch is [row_pos, row_pos + num_rows).

    Status _generate_key_index(std::vector<const IColumn*>& key_columns, const IColumn* seq_column,

                               size_t row_pos, size_t num_rows,

                               const std::map<uint32_t, const IColumn*>& cluster_key_columns);

    Status _generate_primary_key_index(const std::vector<const IColumn*>& primary_key_columns,

                                       const IColumn* seq_column, size_t row_pos, size_t num_rows,

                                       bool need_sort);

    Status _generate_short_key_index(std::vector<const IColumn*>& key_columns, size_t row_pos,

                                     size_t num_rows, const std::vector<size_t>& short_key_pos);

    Status _check_column_writer_disk_capacity(size_t cid);

    Status _finalize_column_writer_and_update_meta(size_t cid);

    bool _is_mow() {

        return _tablet_schema->keys_type() == UNIQUE_KEYS && _opts.enable_unique_key_merge_on_write;

    }

    bool _is_mow_with_cluster_key() {

        return _is_mow() && !_tablet_schema->cluster_key_uids().empty();

    }

private:

    uint32_t _segment_id;

    TabletSchemaSPtr _tablet_schema;

    BaseTabletSPtr _tablet;

    DataDir* _data_dir = nullptr;

    VerticalSegmentWriterOptions _opts;

    // Not owned. owned by RowsetWriter or SegmentFlusher

    io::FileWriter* _file_writer = nullptr;

    // Not owned. owned by RowsetWriter or SegmentFlusher

    IndexFileWriter* _index_file_writer = nullptr;

    SegmentFooterPB _footer;

    SegmentIndexFileCacheInfo _index_file_cache_info;

    size_t _num_short_key_columns;

    std::unique_ptr<ShortKeyIndexBuilder> _short_key_index_builder;

    std::unique_ptr<PrimaryKeyIndexBuilder> _primary_key_index_builder;

    std::vector<std::unique_ptr<ColumnWriter>> _column_writers;

    std::unique_ptr<MemTracker> _mem_tracker;

    // used for building short key index or primary key index during vectorized write.

    // NOTE: must stay declared after _tablet_schema and _opts, the constructor

    // init list reads both through _is_mow().

    RowKeyEncoder _key_encoder;

    size_t _short_key_row_pos = 0;

    // The open column group: the columns it holds, and whether it is the one

    // carrying the key columns. A whole-schema group (init() or write_block) is

    // the degenerate case -- every column, has_key=true.

    std::vector<uint32_t> _column_ids;

    bool _has_key = true;

    // _num_rows_written means row count already written in this current column group

    uint32_t _num_rows_written = 0;

    // _row_count means total row count of this segment

    // In vertical compaction row count is recorded when key columns group finish

    //  and _num_rows_written will be updated in value column group

    uint32_t _row_count = 0;

    bool _is_first_row = true;

    faststring _min_key;

    faststring _max_key;

    // Cluster key MOW only: primary keys collected here until the group's rows

    // are all in, then _sort_primary_keys_into_index() hands them over. The

    // rowset writer rolls a segment on their byte count.

    std::vector<std::string> _primary_keys;

    uint64_t _primary_keys_size = 0;

    // variant statistics calculator for efficient stats collection; only the

    // compaction write type feeds it

    std::unique_ptr<VariantStatsCaculator> _variant_stats_calculator;

    // the derived column the transform chain hands off to write_block's bounded pump

    DerivedColumn _derived_column;

    // true once write_block has written the open group's data pages, so

    // finalize_columns() does not write them again; clear() resets it

    bool _columns_data_flushed = false;

};

} // namespace segment_v2

} // namespace doris