inner_index_parameter.cpp

// Copyright 2024-present the vsag project
//
// Licensed 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 "inner_index_parameter.h"

#include "common.h"
#include "datacell/attribute_inverted_interface_parameter.h"
#include "datacell/extra_info_datacell_parameter.h"
#include "datacell/flatten_datacell_parameter.h"
#include "impl/logger/logger.h"
#include "inner_string_params.h"
#include "utils/param_compat_macros.h"
#include "vsag/constants.h"

namespace vsag {

namespace {

auto
parse_label_remap_type(const std::string& remap_type) -> LabelRemapType {
    if (remap_type == LABEL_REMAP_TYPE_VALUE_PG) {
        return LabelRemapType::PG;
    }
    if (remap_type == LABEL_REMAP_TYPE_VALUE_ROBIN) {
        return LabelRemapType::ROBIN;
    }
    throw VsagException(ErrorType::INVALID_ARGUMENT,
                        fmt::format("invalid label_remap_type: {}", remap_type));
}

auto
dump_label_remap_type(LabelRemapType remap_type) -> const char* {
    return remap_type == LabelRemapType::ROBIN ? LABEL_REMAP_TYPE_VALUE_ROBIN
                                               : LABEL_REMAP_TYPE_VALUE_PG;
}

}  // namespace

void
InnerIndexParameter::FromJson(const JsonType& json) {
    if (json.Contains(USE_REORDER_KEY)) {
        this->use_reorder = json[USE_REORDER_KEY].GetBool();
    }

    if (json.Contains(REORDER_SOURCE_KEY)) {
        this->reorder_source = json[REORDER_SOURCE_KEY].GetString();
    }
    CHECK_ARGUMENT(this->reorder_source == HGRAPH_REORDER_SOURCE_PRECISE ||
                       this->reorder_source == HGRAPH_REORDER_SOURCE_BASE,
                   fmt::format("invalid reorder_source: {}, supported values are \"{}\" and \"{}\"",
                               this->reorder_source,
                               HGRAPH_REORDER_SOURCE_PRECISE,
                               HGRAPH_REORDER_SOURCE_BASE));

    if (json.Contains(USE_ATTRIBUTE_FILTER_KEY)) {
        this->use_attribute_filter = json[USE_ATTRIBUTE_FILTER_KEY].GetBool();
    }

    if (this->use_attribute_filter) {
        this->attr_inverted_interface_param =
            std::make_shared<AttributeInvertedInterfaceParameter>();
        if (json.Contains(ATTR_PARAMS_KEY)) {
            this->attr_inverted_interface_param->FromJson(json[ATTR_PARAMS_KEY]);
        }
    }

    if (json.Contains(BUILD_THREAD_COUNT_KEY)) {
        this->build_thread_count = json[BUILD_THREAD_COUNT_KEY].GetInt();
    }

    if (json.Contains(LABEL_REMAP_TYPE_KEY)) {
        this->label_remap_type = parse_label_remap_type(json[LABEL_REMAP_TYPE_KEY].GetString());
    }

    if (this->use_reorder && this->reorder_source != HGRAPH_REORDER_SOURCE_BASE) {
        CHECK_ARGUMENT(
            json.Contains(PRECISE_CODES_KEY),
            fmt::format("ivf parameters must contains {} when enable reorder", PRECISE_CODES_KEY));
        this->precise_codes_param = CreateFlattenParam(json[PRECISE_CODES_KEY]);
    }

    if (json.Contains(STORE_RAW_VECTOR_KEY)) {
        this->store_raw_vector = json[STORE_RAW_VECTOR_KEY].GetBool();
    }

    if (this->store_raw_vector) {
        this->raw_vector_param = CreateFlattenParam(json[RAW_VECTOR_KEY]);
    }

    if (json.Contains(EXTRA_INFO_KEY)) {
        this->extra_info_param = std::make_shared<ExtraInfoDataCellParameter>();
        this->extra_info_param->FromJson(json[EXTRA_INFO_KEY]);
    }

    if (json.Contains(TRAIN_SAMPLE_COUNT_KEY)) {
        this->train_sample_count = json[TRAIN_SAMPLE_COUNT_KEY].GetInt();
        CHECK_ARGUMENT(
            this->train_sample_count >= 512,
            fmt::format("train_sample_count must be greater than or equal to 512, got: {}",
                        this->train_sample_count));
        CHECK_ARGUMENT(
            this->train_sample_count <= 65536L,
            fmt::format("train_sample_count must be less than or equal to 65536, got: {}",
                        this->train_sample_count));
    }
}

JsonType
InnerIndexParameter::ToJson() const {
    JsonType json;
    json[USE_REORDER_KEY].SetBool(this->use_reorder);
    json[REORDER_SOURCE_KEY].SetString(this->reorder_source);
    json[BUILD_THREAD_COUNT_KEY].SetInt(this->build_thread_count);
    json[LABEL_REMAP_TYPE_KEY].SetString(dump_label_remap_type(this->label_remap_type));
    json[USE_ATTRIBUTE_FILTER_KEY].SetBool(this->use_attribute_filter);
    if (use_reorder && this->reorder_source != HGRAPH_REORDER_SOURCE_BASE) {
        json[PRECISE_CODES_KEY].SetJson(this->precise_codes_param->ToJson());
    }
    json[STORE_RAW_VECTOR_KEY].SetBool(this->store_raw_vector);
    if (this->store_raw_vector) {
        json[RAW_VECTOR_KEY].SetJson(this->raw_vector_param->ToJson());
    }
    if (this->extra_info_param) {
        json[EXTRA_INFO_KEY].SetJson(this->extra_info_param->ToJson());
    }
    if (this->use_attribute_filter) {
        json[ATTR_PARAMS_KEY].SetJson(this->attr_inverted_interface_param->ToJson());
    }

    json[TRAIN_SAMPLE_COUNT_KEY].SetInt(this->train_sample_count);

    auto str = json.Dump(4);
    return json;
}
bool
InnerIndexParameter::CheckCompatibility(const ParamPtr& other) const {
    PARAM_CAST_OR_RETURN(InnerIndexParameter, p, other);
    CHECK_FIELD_EQ(*this, *p, use_reorder);
    CHECK_FIELD_EQ(*this, *p, reorder_source);
    if (this->use_reorder && this->reorder_source != HGRAPH_REORDER_SOURCE_BASE) {
        CHECK_SUB_PARAM(*this, *p, precise_codes_param);
    }
    CHECK_FIELD_EQ(*this, *p, use_attribute_filter);
    CHECK_FIELD_EQ(*this, *p, label_remap_type);
    return true;
}
}  // namespace vsag