be/src/storage/index/indexed_column_writer.cpp

// 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/index/indexed_column_writer.h"

#include <gen_cpp/segment_v2.pb.h>

#include <ostream>
#include <string>

#include "common/logging.h"
#include "io/fs/file_writer.h"
#include "storage/index/index_page.h"
#include "storage/key_coder.h"
#include "storage/olap_common.h"
#include "storage/segment/encoding_info.h"
#include "storage/segment/options.h"
#include "storage/segment/page_builder.h"
#include "storage/segment/page_io.h"
#include "storage/segment/page_pointer.h"
#include "storage/types.h"
#include "util/block_compression.h"
#include "util/slice.h"

namespace doris {
namespace segment_v2 {

IndexedColumnWriter::IndexedColumnWriter(const IndexedColumnWriterOptions& options, FieldType type,
                                         io::FileWriter* file_writer)
        : _options(options),
          _type(type),
          _file_writer(file_writer),
          _num_values(0),
          _num_data_pages(0),
          _disk_size(0),
          _value_key_coder(nullptr),
          _compress_codec(nullptr) {}

IndexedColumnWriter::~IndexedColumnWriter() = default;

Status IndexedColumnWriter::init() {
    // Caller must set _options.encoding to a concrete value before calling init.
    if (_options.encoding == DEFAULT_ENCODING) {
        return Status::InternalError(
                "IndexedColumnWriterOptions::encoding is DEFAULT_ENCODING for type={}; caller must "
                "resolve to a concrete encoding before IndexedColumnWriter::init",
                _type);
    }
    const EncodingInfo* encoding_info;
    RETURN_IF_ERROR(EncodingInfo::get(_type, _options.encoding, &encoding_info));

    PageBuilder* data_page_builder = nullptr;
    PageBuilderOptions builder_option;
    builder_option.need_check_bitmap = false;
    builder_option.data_page_size = _options.data_page_size;
    RETURN_IF_ERROR(encoding_info->create_page_builder(builder_option, &data_page_builder));
    _data_page_builder.reset(data_page_builder);

    if (_options.write_ordinal_index) {
        _ordinal_index_builder.reset(new IndexPageBuilder(_options.index_page_size, true));
    }
    if (_options.write_value_index) {
        _value_index_builder.reset(new IndexPageBuilder(_options.index_page_size, true));
        _value_key_coder = get_key_coder(_type);
    }

    if (_options.compression != NO_COMPRESSION) {
        RETURN_IF_ERROR(get_block_compression_codec(_options.compression, &_compress_codec));
    }
    return Status::OK();
}

Status IndexedColumnWriter::add(const void* value) {
    if (_options.write_value_index && _data_page_builder->count() == 0) {
        // remember page's first value encoded key because it's used to build value index
        _first_value_string.clear();
        _value_key_coder->full_encode_ascending(value, &_first_value_string);
    }
    size_t num_to_write = 1;
    RETURN_IF_ERROR(
            _data_page_builder->add(reinterpret_cast<const uint8_t*>(value), &num_to_write));
    CHECK(num_to_write == 1 || num_to_write == 0);
    if (num_to_write == 0) {
        CHECK(_data_page_builder->is_page_full());
        // current page is already full, we need to first flush the current page,
        // and then add the value to the new page
        size_t num_val;
        RETURN_IF_ERROR(_finish_current_data_page(num_val));
        return add(value);
    }
    _num_values++;
    size_t num_val;
    if (_data_page_builder->is_page_full()) {
        RETURN_IF_ERROR(_finish_current_data_page(num_val));
    }
    return Status::OK();
}

Status IndexedColumnWriter::_finish_current_data_page(size_t& num_val) {
    auto num_values_in_page = _data_page_builder->count();
    num_val = num_values_in_page;
    if (num_values_in_page == 0) {
        return Status::OK();
    }
    ordinal_t first_ordinal = _num_values - num_values_in_page;

    // IndexedColumn doesn't have NULLs, thus data page body only contains encoded values
    OwnedSlice page_body;
    RETURN_IF_ERROR(_data_page_builder->finish(&page_body));
    RETURN_IF_ERROR(_data_page_builder->reset());

    PageFooterPB footer;
    footer.set_type(DATA_PAGE);
    footer.set_uncompressed_size(static_cast<uint32_t>(page_body.slice().get_size()));
    footer.mutable_data_page_footer()->set_first_ordinal(first_ordinal);
    footer.mutable_data_page_footer()->set_num_values(num_values_in_page);
    footer.mutable_data_page_footer()->set_nullmap_size(0);

    uint64_t start_size = _file_writer->bytes_appended();
    RETURN_IF_ERROR(PageIO::compress_and_write_page(
            _compress_codec, _options.compression_min_space_saving, _file_writer,
            {page_body.slice()}, footer, &_last_data_page));
    _num_data_pages++;
    _disk_size += (_file_writer->bytes_appended() - start_size);

    if (_options.write_ordinal_index) {
        std::string key;
        KeyCoderTraits<FieldType::OLAP_FIELD_TYPE_UNSIGNED_BIGINT>::full_encode_ascending(
                &first_ordinal, &key);
        _ordinal_index_builder->add(key, _last_data_page);
    }

    if (_options.write_value_index) {
        // TODO short separate key optimize
        _value_index_builder->add(_first_value_string, _last_data_page);
        // TODO record last key in short separate key optimize
    }
    return Status::OK();
}

Status IndexedColumnWriter::finish(IndexedColumnMetaPB* meta) {
    size_t num_val_in_page;
    RETURN_IF_ERROR(_finish_current_data_page(num_val_in_page));
    if (_options.write_ordinal_index) {
        RETURN_IF_ERROR(
                _flush_index(_ordinal_index_builder.get(), meta->mutable_ordinal_index_meta()));
    }
    if (_options.write_value_index) {
        RETURN_IF_ERROR(_flush_index(_value_index_builder.get(), meta->mutable_value_index_meta()));
    }
    meta->set_data_type(int(_type));
    meta->set_encoding(_options.encoding);
    meta->set_num_values(_num_values);
    meta->set_compression(_options.compression);
    // `_finish_current_data_page` will be called in `add` function when page is full,
    // so num_val_in_page will be zero in this case.
    if (_num_data_pages <= 1 && num_val_in_page != 0) {
        DCHECK(num_val_in_page == _num_values)
                << "num_val_in_page: " << num_val_in_page << ", _num_values: " << _num_values;
    }
    return Status::OK();
}

Status IndexedColumnWriter::_flush_index(IndexPageBuilder* index_builder, BTreeMetaPB* meta) {
    if (_num_data_pages <= 1) {
        meta->set_is_root_data_page(true);
        _last_data_page.to_proto(meta->mutable_root_page());
    } else {
        OwnedSlice page_body;
        PageFooterPB page_footer;
        index_builder->finish(&page_body, &page_footer);

        PagePointer pp;
        uint64_t start_size = _file_writer->bytes_appended();
        RETURN_IF_ERROR(PageIO::compress_and_write_page(
                _compress_codec, _options.compression_min_space_saving, _file_writer,
                {page_body.slice()}, page_footer, &pp));
        _disk_size += (_file_writer->bytes_appended() - start_size);

        meta->set_is_root_data_page(false);
        pp.to_proto(meta->mutable_root_page());
    }
    return Status::OK();
}

} // namespace segment_v2
} // namespace doris