// 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/rowset/segment_creator.h"

// IWYU pragma: no_include <bthread/errno.h>

#include <cerrno> // IWYU pragma: keep

#include <filesystem>

#include <memory>

#include <sstream>

#include <utility>

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

#include "common/config.h"

#include "common/exception.h"

#include "common/logging.h"

#include "common/status.h"

#include "core/assert_cast.h"

#include "core/block/block.h"

#include "core/block/columns_with_type_and_name.h"

#include "core/column/column.h"

#include "core/column/column_nullable.h"

#include "core/column/column_string.h"

#include "core/column/column_variant.h"

#include "core/data_type/data_type.h"

#include "core/types.h"

#include "io/fs/file_writer.h"

#include "storage/olap_define.h"

#include "storage/rowset/beta_rowset_writer.h" // SegmentStatistics

#include "storage/segment/segment_writer.h"

#include "storage/segment/vertical_segment_writer.h"

#include "storage/tablet/tablet_schema.h"

#include "storage/utils.h"

#include "util/json/json_parser.h"

#include "util/pretty_printer.h"

#include "util/stopwatch.hpp"

namespace doris {

using namespace ErrorCode;

SegmentFlusher::SegmentFlusher(RowsetWriterContext& context, SegmentFileCollection& seg_files,

                               InvertedIndexFileCollection& idx_files)

        : _context(context), _seg_files(seg_files), _idx_files(idx_files) {}

SegmentFlusher::~SegmentFlusher() = default;

// NOLINTNEXTLINE(readability-function-cognitive-complexity)

Status SegmentFlusher::flush_single_block(const Block* block, int32_t segment_id,

                                          int64_t* flush_size) {

    if (block->rows() == 0) {

        return Status::OK();

    }

    Block flush_block(*block);

    bool no_compression = flush_block.bytes() <= config::segment_compression_threshold_kb * 1024;

    if (config::enable_vertical_segment_writer) {

        std::unique_ptr<segment_v2::VerticalSegmentWriter> writer;

        RETURN_IF_ERROR(_create_segment_writer(writer, segment_id, no_compression));

        RETURN_IF_ERROR_OR_CATCH_EXCEPTION(_add_rows(writer, &flush_block, 0, flush_block.rows()));

        RETURN_IF_ERROR(_flush_segment_writer(writer, flush_size));

    } else {

        std::unique_ptr<segment_v2::SegmentWriter> writer;

        RETURN_IF_ERROR(_create_segment_writer(writer, segment_id, no_compression));

        RETURN_IF_ERROR_OR_CATCH_EXCEPTION(_add_rows(writer, &flush_block, 0, flush_block.rows()));

        RETURN_IF_ERROR(_flush_segment_writer(writer, flush_size));

    }

    return Status::OK();

}

Status SegmentFlusher::close() {

    RETURN_IF_ERROR(_seg_files.close());

    RETURN_IF_ERROR(_idx_files.finish_close());

    return Status::OK();

}

Status SegmentFlusher::_add_rows(std::unique_ptr<segment_v2::SegmentWriter>& segment_writer,

                                 const Block* block, size_t row_pos, size_t num_rows) {

    RETURN_IF_ERROR(segment_writer->append_block(block, row_pos, num_rows));

    _num_rows_written += num_rows;

    return Status::OK();

}

Status SegmentFlusher::_add_rows(std::unique_ptr<segment_v2::VerticalSegmentWriter>& segment_writer,

                                 const Block* block, size_t row_pos, size_t num_rows) {

    RETURN_IF_ERROR(segment_writer->batch_block(block, row_pos, num_rows));

    RETURN_IF_ERROR(segment_writer->write_batch());

    _num_rows_written += num_rows;

    return Status::OK();

}

Status SegmentFlusher::_create_segment_writer(std::unique_ptr<segment_v2::SegmentWriter>& writer,

                                              int32_t segment_id, bool no_compression) {

    io::FileWriterPtr segment_file_writer;

    RETURN_IF_ERROR(_context.file_writer_creator->create(segment_id, segment_file_writer));

    IndexFileWriterPtr index_file_writer;

    if (_context.tablet_schema->has_inverted_index() || _context.tablet_schema->has_ann_index()) {

        RETURN_IF_ERROR(_context.file_writer_creator->create(segment_id, &index_file_writer));

    }

    segment_v2::SegmentWriterOptions writer_options;

    writer_options.enable_unique_key_merge_on_write = _context.enable_unique_key_merge_on_write;

    writer_options.rowset_ctx = &_context;

    writer_options.write_type = _context.write_type;

    writer_options.max_rows_per_segment = _context.max_rows_per_segment;

    writer_options.mow_ctx = _context.mow_context;

    if (no_compression) {

        writer_options.compression_type = NO_COMPRESSION;

    }

    writer = std::make_unique<segment_v2::SegmentWriter>(

            segment_file_writer.get(), segment_id, _context.tablet_schema, _context.tablet,

            _context.data_dir, writer_options, index_file_writer.get());

    RETURN_IF_ERROR(_seg_files.add(segment_id, std::move(segment_file_writer)));

    if (_context.tablet_schema->has_inverted_index() || _context.tablet_schema->has_ann_index()) {

        RETURN_IF_ERROR(_idx_files.add(segment_id, std::move(index_file_writer)));

    }

    auto s = writer->init();

    if (!s.ok()) {

        LOG(WARNING) << "failed to init segment writer: " << s.to_string();

        writer.reset();

        return s;

    }

    return Status::OK();

}

Status SegmentFlusher::_create_segment_writer(

        std::unique_ptr<segment_v2::VerticalSegmentWriter>& writer, int32_t segment_id,

        bool no_compression) {

    io::FileWriterPtr segment_file_writer;

    RETURN_IF_ERROR(_context.file_writer_creator->create(segment_id, segment_file_writer));

    IndexFileWriterPtr index_file_writer;

    if (_context.tablet_schema->has_inverted_index() || _context.tablet_schema->has_ann_index()) {

        RETURN_IF_ERROR(_context.file_writer_creator->create(segment_id, &index_file_writer));

    }

    segment_v2::VerticalSegmentWriterOptions writer_options;

    writer_options.enable_unique_key_merge_on_write = _context.enable_unique_key_merge_on_write;

    writer_options.rowset_ctx = &_context;

    writer_options.write_type = _context.write_type;

    writer_options.mow_ctx = _context.mow_context;

    if (no_compression) {

        writer_options.compression_type = NO_COMPRESSION;

    }

    writer = std::make_unique<segment_v2::VerticalSegmentWriter>(

            segment_file_writer.get(), segment_id, _context.tablet_schema, _context.tablet,

            _context.data_dir, writer_options, index_file_writer.get());

    RETURN_IF_ERROR(_seg_files.add(segment_id, std::move(segment_file_writer)));

    if (_context.tablet_schema->has_inverted_index() || _context.tablet_schema->has_ann_index()) {

        RETURN_IF_ERROR(_idx_files.add(segment_id, std::move(index_file_writer)));

    }

    auto s = writer->init();

    if (!s.ok()) {

        LOG(WARNING) << "failed to init segment writer: " << s.to_string();

        writer.reset();

        return s;

    }

    VLOG_DEBUG << "create new segment writer, tablet_id:" << _context.tablet_id

               << " segment id: " << segment_id << " filename: " << writer->data_dir_path()

               << " rowset_id:" << _context.rowset_id;

    return Status::OK();

}

Status SegmentFlusher::_flush_segment_writer(

        std::unique_ptr<segment_v2::VerticalSegmentWriter>& writer, int64_t* flush_size) {

    MonotonicStopWatch total_timer;

    total_timer.start();

    uint32_t row_num = writer->num_rows_written();

    _num_rows_updated += writer->num_rows_updated();

    _num_rows_deleted += writer->num_rows_deleted();

    _num_rows_new_added += writer->num_rows_new_added();

    _num_rows_filtered += writer->num_rows_filtered();

    if (row_num == 0) {

        return Status::OK();

    }

    MonotonicStopWatch finalize_timer;

    finalize_timer.start();

    uint64_t segment_file_size;

    uint64_t common_index_size;

    Status s = writer->finalize(&segment_file_size, &common_index_size);

    finalize_timer.stop();

    if (!s.ok()) {

        return Status::Error(s.code(), "failed to finalize segment: {}", s.to_string());

    }

    MonotonicStopWatch inverted_index_timer;

    inverted_index_timer.start();

    int64_t inverted_index_file_size = 0;

    RETURN_IF_ERROR(writer->close_inverted_index(&inverted_index_file_size));

    inverted_index_timer.stop();

    VLOG_DEBUG << "tablet_id:" << _context.tablet_id

               << " flushing filename: " << writer->data_dir_path()

               << " rowset_id:" << _context.rowset_id;

    KeyBoundsPB key_bounds;

    Slice min_key = writer->min_encoded_key();

    Slice max_key = writer->max_encoded_key();

    DCHECK_LE(min_key.compare(max_key), 0);

    key_bounds.set_min_key(min_key.to_string());

    key_bounds.set_max_key(max_key.to_string());

    uint32_t segment_id = writer->segment_id();

    SegmentStatistics segstat;

    segstat.row_num = row_num;

    segstat.data_size = segment_file_size;

    segstat.index_size = inverted_index_file_size;

    segstat.key_bounds = key_bounds;

    writer.reset();

    MonotonicStopWatch collector_timer;

    collector_timer.start();

    RETURN_IF_ERROR(_context.segment_collector->add(segment_id, segstat));

    collector_timer.stop();

    total_timer.stop();

    LOG(INFO) << "tablet_id:" << _context.tablet_id

              << ", flushing rowset_dir: " << _context.tablet_path

              << ", rowset_id:" << _context.rowset_id

              << ", data size:" << PrettyPrinter::print_bytes(segstat.data_size)

              << ", index size:" << PrettyPrinter::print_bytes(segstat.index_size)

              << ", timing breakdown: total=" << total_timer.elapsed_time_milliseconds() << "ms"

              << ", finalize=" << finalize_timer.elapsed_time_milliseconds() << "ms"

              << ", inverted_index=" << inverted_index_timer.elapsed_time_milliseconds() << "ms"

              << ", collector=" << collector_timer.elapsed_time_milliseconds() << "ms";

    if (flush_size) {

        *flush_size = segment_file_size;

    }

    return Status::OK();

}

Status SegmentFlusher::_flush_segment_writer(std::unique_ptr<segment_v2::SegmentWriter>& writer,

                                             int64_t* flush_size) {

    MonotonicStopWatch total_timer;

    total_timer.start();

    uint32_t row_num = writer->num_rows_written();

    _num_rows_updated += writer->num_rows_updated();

    _num_rows_deleted += writer->num_rows_deleted();

    _num_rows_new_added += writer->num_rows_new_added();

    _num_rows_filtered += writer->num_rows_filtered();

    if (row_num == 0) {

        return Status::OK();

    }

    MonotonicStopWatch finalize_timer;

    finalize_timer.start();

    uint64_t segment_file_size;

    uint64_t common_index_size;

    Status s = writer->finalize(&segment_file_size, &common_index_size);

    finalize_timer.stop();

    if (!s.ok()) {

        return Status::Error(s.code(), "failed to finalize segment: {}", s.to_string());

    }

    MonotonicStopWatch inverted_index_timer;

    inverted_index_timer.start();

    int64_t inverted_index_file_size = 0;

    RETURN_IF_ERROR(writer->close_inverted_index(&inverted_index_file_size));

    inverted_index_timer.stop();

    VLOG_DEBUG << "tablet_id:" << _context.tablet_id

               << " flushing rowset_dir: " << _context.tablet_path

               << " rowset_id:" << _context.rowset_id;

    KeyBoundsPB key_bounds;

    Slice min_key = writer->min_encoded_key();

    Slice max_key = writer->max_encoded_key();

    DCHECK_LE(min_key.compare(max_key), 0);

    key_bounds.set_min_key(min_key.to_string());

    key_bounds.set_max_key(max_key.to_string());

    uint32_t segment_id = writer->get_segment_id();

    SegmentStatistics segstat;

    segstat.row_num = row_num;

    segstat.data_size = segment_file_size;

    segstat.index_size = inverted_index_file_size;

    segstat.key_bounds = key_bounds;

    writer.reset();

    MonotonicStopWatch collector_timer;

    collector_timer.start();

    RETURN_IF_ERROR(_context.segment_collector->add(segment_id, segstat));

    collector_timer.stop();

    total_timer.stop();

    LOG(INFO) << "tablet_id:" << _context.tablet_id

              << ", flushing rowset_dir: " << _context.tablet_path

              << ", rowset_id:" << _context.rowset_id

              << ", data size:" << PrettyPrinter::print_bytes(segstat.data_size)

              << ", index size:" << PrettyPrinter::print_bytes(segstat.index_size)

              << ", timing breakdown: total=" << total_timer.elapsed_time_milliseconds() << "ms"

              << ", finalize=" << finalize_timer.elapsed_time_milliseconds() << "ms"

              << ", inverted_index=" << inverted_index_timer.elapsed_time_milliseconds() << "ms"

              << ", collector=" << collector_timer.elapsed_time_milliseconds() << "ms";

    if (flush_size) {

        *flush_size = segment_file_size;

    }

    return Status::OK();

}

Status SegmentFlusher::create_writer(std::unique_ptr<SegmentFlusher::Writer>& writer,

                                     uint32_t segment_id) {

    std::unique_ptr<segment_v2::SegmentWriter> segment_writer;

    RETURN_IF_ERROR(_create_segment_writer(segment_writer, segment_id));

    DCHECK(segment_writer != nullptr);

    writer.reset(new SegmentFlusher::Writer(this, segment_writer));

    return Status::OK();

}

SegmentFlusher::Writer::Writer(SegmentFlusher* flusher,

                               std::unique_ptr<segment_v2::SegmentWriter>& segment_writer)

        : _flusher(flusher), _writer(std::move(segment_writer)) {};

SegmentFlusher::Writer::~Writer() = default;

Status SegmentFlusher::Writer::flush() {

    return _flusher->_flush_segment_writer(_writer);

}

int64_t SegmentFlusher::Writer::max_row_to_add(size_t row_avg_size_in_bytes) {

    return _writer->max_row_to_add(row_avg_size_in_bytes);

}

SegmentCreator::SegmentCreator(RowsetWriterContext& context, SegmentFileCollection& seg_files,

                               InvertedIndexFileCollection& idx_files)

        : _segment_flusher(context, seg_files, idx_files) {}

Status SegmentCreator::add_block(const Block* block) {

    if (block->rows() == 0) {

        return Status::OK();

    }

    size_t block_size_in_bytes = block->bytes();

    size_t block_row_num = block->rows();

    size_t row_avg_size_in_bytes = std::max((size_t)1, block_size_in_bytes / block_row_num);

    size_t row_offset = 0;

    if (_flush_writer == nullptr) {

        RETURN_IF_ERROR(_segment_flusher.create_writer(_flush_writer, allocate_segment_id()));

    }

    do {

        auto max_row_add = _flush_writer->max_row_to_add(row_avg_size_in_bytes);

        if (UNLIKELY(max_row_add < 1)) {

            // no space for another single row, need flush now

            RETURN_IF_ERROR(flush());

            RETURN_IF_ERROR(_segment_flusher.create_writer(_flush_writer, allocate_segment_id()));

            max_row_add = _flush_writer->max_row_to_add(row_avg_size_in_bytes);

            DCHECK(max_row_add > 0);

        }

        size_t input_row_num = std::min(block_row_num - row_offset, size_t(max_row_add));

        RETURN_IF_ERROR(_flush_writer->add_rows(block, row_offset, input_row_num));

        row_offset += input_row_num;

    } while (row_offset < block_row_num);

    return Status::OK();

}

Status SegmentCreator::flush() {

    if (_flush_writer == nullptr) {

        return Status::OK();

    }

    RETURN_IF_ERROR(_flush_writer->flush());

    _flush_writer.reset();

    return Status::OK();

}

Status SegmentCreator::flush_single_block(const Block* block, int32_t segment_id,

                                          int64_t* flush_size) {

    if (block->rows() == 0) {

        return Status::OK();

    }

    RETURN_IF_ERROR(_segment_flusher.flush_single_block(block, segment_id, flush_size));

    return Status::OK();

}

Status SegmentCreator::close() {

    RETURN_IF_ERROR(flush());

    RETURN_IF_ERROR(_segment_flusher.close());

    return Status::OK();

}

} // namespace doris