281 lines
9.4 KiB
C++
281 lines
9.4 KiB
C++
// Copyright 2025-present the zvec 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 "zvec/db/index_params.h"
|
|
#include <gtest/gtest.h>
|
|
|
|
using namespace zvec;
|
|
|
|
TEST(IndexParamsTest, IndexParamsBaseClass) {
|
|
// Test that IndexParams is abstract and can't be instantiated directly
|
|
// This is more of a compile-time check - we can't directly instantiate an
|
|
// abstract class
|
|
|
|
// Test is_vector_index_type method
|
|
HnswIndexParams hnsw_params(MetricType::L2, 16, 100);
|
|
EXPECT_TRUE(hnsw_params.is_vector_index_type());
|
|
|
|
FlatIndexParams flat_params(MetricType::IP);
|
|
EXPECT_TRUE(flat_params.is_vector_index_type());
|
|
|
|
IVFIndexParams ivf_params(MetricType::COSINE, 100);
|
|
EXPECT_TRUE(ivf_params.is_vector_index_type());
|
|
|
|
InvertIndexParams invert_params(true);
|
|
EXPECT_FALSE(invert_params.is_vector_index_type());
|
|
}
|
|
|
|
TEST(IndexParamsTest, InvertIndexParams) {
|
|
// Test constructor
|
|
InvertIndexParams params(true);
|
|
EXPECT_EQ(params.type(), IndexType::INVERT);
|
|
EXPECT_TRUE(params.enable_range_optimization());
|
|
|
|
InvertIndexParams params2(false);
|
|
EXPECT_FALSE(params2.enable_range_optimization());
|
|
|
|
// Test clone method
|
|
auto cloned = params.clone();
|
|
EXPECT_NE(cloned.get(), ¶ms); // Should be different objects
|
|
EXPECT_EQ(cloned->type(), IndexType::INVERT);
|
|
|
|
// Test comparison operators
|
|
InvertIndexParams params3(true);
|
|
InvertIndexParams params4(false);
|
|
|
|
EXPECT_TRUE(params == params3);
|
|
EXPECT_FALSE(params == params4);
|
|
EXPECT_TRUE(params != params4);
|
|
|
|
// Test setter
|
|
params2.set_enable_range_optimization(true);
|
|
EXPECT_TRUE(params2.enable_range_optimization());
|
|
EXPECT_TRUE(params2 == params);
|
|
}
|
|
|
|
TEST(IndexParamsTest, VectorIndexParamsBase) {
|
|
// Test constructor and basic methods
|
|
FlatIndexParams flat_params(MetricType::L2, QuantizeType::FP16);
|
|
EXPECT_EQ(flat_params.type(), IndexType::FLAT);
|
|
EXPECT_EQ(flat_params.metric_type(), MetricType::L2);
|
|
EXPECT_EQ(flat_params.quantize_type(), QuantizeType::FP16);
|
|
|
|
// Test setters
|
|
flat_params.set_metric_type(MetricType::IP);
|
|
EXPECT_EQ(flat_params.metric_type(), MetricType::IP);
|
|
|
|
flat_params.set_quantize_type(QuantizeType::INT8);
|
|
EXPECT_EQ(flat_params.quantize_type(), QuantizeType::INT8);
|
|
}
|
|
|
|
TEST(IndexParamsTest, HnswIndexParams) {
|
|
// Test constructor
|
|
HnswIndexParams params(MetricType::COSINE, 20, 150, QuantizeType::INT4);
|
|
EXPECT_EQ(params.type(), IndexType::HNSW);
|
|
EXPECT_EQ(params.metric_type(), MetricType::COSINE);
|
|
EXPECT_EQ(params.m(), 20);
|
|
EXPECT_EQ(params.ef_construction(), 150);
|
|
EXPECT_EQ(params.quantize_type(), QuantizeType::INT4);
|
|
|
|
// Test clone
|
|
auto cloned = params.clone();
|
|
EXPECT_NE(cloned.get(), ¶ms);
|
|
EXPECT_EQ(cloned->type(), IndexType::HNSW);
|
|
|
|
// Test comparison
|
|
HnswIndexParams params2(MetricType::COSINE, 20, 150, QuantizeType::INT4);
|
|
HnswIndexParams params3(MetricType::L2, 20, 150, QuantizeType::INT4);
|
|
HnswIndexParams params4(MetricType::COSINE, 16, 150, QuantizeType::INT4);
|
|
HnswIndexParams params5(MetricType::COSINE, 20, 200, QuantizeType::INT4);
|
|
|
|
EXPECT_TRUE(params == params2);
|
|
EXPECT_FALSE(params == params3);
|
|
EXPECT_FALSE(params == params4);
|
|
EXPECT_FALSE(params == params5);
|
|
|
|
// Test setters
|
|
params.set_m(10);
|
|
EXPECT_EQ(params.m(), 10);
|
|
|
|
params.set_ef_construction(75);
|
|
EXPECT_EQ(params.ef_construction(), 75);
|
|
}
|
|
|
|
TEST(IndexParamsTest, FlatIndexParams) {
|
|
// Test constructor
|
|
FlatIndexParams params(MetricType::IP, QuantizeType::FP16);
|
|
EXPECT_EQ(params.type(), IndexType::FLAT);
|
|
EXPECT_EQ(params.metric_type(), MetricType::IP);
|
|
EXPECT_EQ(params.quantize_type(), QuantizeType::FP16);
|
|
|
|
// Test clone
|
|
auto cloned = params.clone();
|
|
EXPECT_NE(cloned.get(), ¶ms);
|
|
EXPECT_EQ(cloned->type(), IndexType::FLAT);
|
|
|
|
// Test comparison
|
|
FlatIndexParams params2(MetricType::IP, QuantizeType::FP16);
|
|
FlatIndexParams params3(MetricType::L2, QuantizeType::FP16);
|
|
FlatIndexParams params4(MetricType::IP, QuantizeType::INT8);
|
|
|
|
EXPECT_TRUE(params == params2);
|
|
EXPECT_FALSE(params == params3);
|
|
EXPECT_FALSE(params == params4);
|
|
}
|
|
|
|
TEST(IndexParamsTest, IVFIndexParams) {
|
|
// Test constructor
|
|
IVFIndexParams params(MetricType::L2, 128, 10, false, QuantizeType::INT8);
|
|
EXPECT_EQ(params.type(), IndexType::IVF);
|
|
EXPECT_EQ(params.metric_type(), MetricType::L2);
|
|
EXPECT_EQ(params.n_list(), 128);
|
|
EXPECT_EQ(params.quantize_type(), QuantizeType::INT8);
|
|
|
|
// Test clone
|
|
auto cloned = params.clone();
|
|
EXPECT_NE(cloned.get(), ¶ms);
|
|
EXPECT_EQ(cloned->type(), IndexType::IVF);
|
|
|
|
// Test comparison
|
|
IVFIndexParams params2(MetricType::L2, 128, 10, false, QuantizeType::INT8);
|
|
IVFIndexParams params3(MetricType::IP, 128, 10, false, QuantizeType::INT8);
|
|
IVFIndexParams params4(MetricType::L2, 256, 10, false, QuantizeType::INT8);
|
|
IVFIndexParams params5(MetricType::L2, 128, 10, false, QuantizeType::FP16);
|
|
|
|
EXPECT_TRUE(params == params2);
|
|
EXPECT_FALSE(params == params3);
|
|
EXPECT_FALSE(params == params4);
|
|
EXPECT_FALSE(params == params5);
|
|
|
|
// Test setter
|
|
params.set_n_list(64);
|
|
EXPECT_EQ(params.n_list(), 64);
|
|
}
|
|
|
|
TEST(IndexParamsTest, DefaultVectorIndexParams) {
|
|
// Test default vector index params
|
|
EXPECT_EQ(DefaultVectorIndexParams.type(), IndexType::FLAT);
|
|
EXPECT_EQ(DefaultVectorIndexParams.metric_type(), MetricType::IP);
|
|
EXPECT_EQ(DefaultVectorIndexParams.quantize_type(), QuantizeType::UNDEFINED);
|
|
}
|
|
|
|
TEST(IndexParamsTest, DynamicPointerCast) {
|
|
// Test dynamic_pointer_cast functionality with IndexParams
|
|
IndexParams::Ptr base_ptr =
|
|
std::make_shared<HnswIndexParams>(MetricType::L2, 16, 100);
|
|
auto hnsw_ptr = std::dynamic_pointer_cast<HnswIndexParams>(base_ptr);
|
|
EXPECT_NE(hnsw_ptr, nullptr);
|
|
EXPECT_EQ(hnsw_ptr->type(), IndexType::HNSW);
|
|
|
|
// Test casting to wrong type
|
|
auto flat_ptr = std::dynamic_pointer_cast<FlatIndexParams>(base_ptr);
|
|
EXPECT_EQ(flat_ptr, nullptr);
|
|
|
|
// Test casting from base class reference
|
|
IndexParams &base_ref = *base_ptr;
|
|
auto &hnsw_ref = dynamic_cast<HnswIndexParams &>(base_ref);
|
|
EXPECT_EQ(hnsw_ref.type(), IndexType::HNSW);
|
|
}
|
|
|
|
// ==================== QuantizerParam tests ====================
|
|
|
|
TEST(IndexParamsTest, QuantizerParamBasic) {
|
|
// Default constructor: enable_rotate should be false
|
|
QuantizerParam qp_default;
|
|
EXPECT_FALSE(qp_default.enable_rotate());
|
|
|
|
// Constructor with true
|
|
QuantizerParam qp_true(true);
|
|
EXPECT_TRUE(qp_true.enable_rotate());
|
|
|
|
// Constructor with false
|
|
QuantizerParam qp_false(false);
|
|
EXPECT_FALSE(qp_false.enable_rotate());
|
|
|
|
// Setter
|
|
qp_default.set_enable_rotate(true);
|
|
EXPECT_TRUE(qp_default.enable_rotate());
|
|
qp_default.set_enable_rotate(false);
|
|
EXPECT_FALSE(qp_default.enable_rotate());
|
|
|
|
// Equality
|
|
EXPECT_TRUE(qp_true == QuantizerParam(true));
|
|
EXPECT_TRUE(qp_false == QuantizerParam(false));
|
|
EXPECT_FALSE(qp_true == qp_false);
|
|
|
|
// Inequality
|
|
EXPECT_TRUE(qp_true != qp_false);
|
|
EXPECT_FALSE(qp_true != QuantizerParam(true));
|
|
}
|
|
|
|
TEST(IndexParamsTest, QuantizerParamWithVectorIndex) {
|
|
// HnswIndexParams
|
|
{
|
|
HnswIndexParams params(MetricType::COSINE, 16, 100, QuantizeType::INT8);
|
|
EXPECT_FALSE(params.quantizer_param().enable_rotate());
|
|
EXPECT_FALSE(params.enable_rotate()); // convenience getter
|
|
|
|
params.set_quantizer_param(QuantizerParam(true));
|
|
EXPECT_TRUE(params.quantizer_param().enable_rotate());
|
|
EXPECT_TRUE(params.enable_rotate());
|
|
|
|
// Clone preserves quantizer_param
|
|
auto cloned = params.clone();
|
|
auto *cloned_hnsw = dynamic_cast<HnswIndexParams *>(cloned.get());
|
|
ASSERT_NE(cloned_hnsw, nullptr);
|
|
EXPECT_TRUE(cloned_hnsw->quantizer_param().enable_rotate());
|
|
EXPECT_TRUE(*cloned == params);
|
|
|
|
// Equality: different enable_rotate -> not equal
|
|
HnswIndexParams params2(MetricType::COSINE, 16, 100, QuantizeType::INT8);
|
|
params2.set_quantizer_param(QuantizerParam(false));
|
|
EXPECT_FALSE(params == params2);
|
|
}
|
|
|
|
// FlatIndexParams
|
|
{
|
|
FlatIndexParams params(MetricType::L2, QuantizeType::INT8);
|
|
EXPECT_FALSE(params.quantizer_param().enable_rotate());
|
|
|
|
params.set_quantizer_param(QuantizerParam(true));
|
|
EXPECT_TRUE(params.quantizer_param().enable_rotate());
|
|
EXPECT_TRUE(params.enable_rotate());
|
|
|
|
auto cloned = params.clone();
|
|
auto *cloned_flat = dynamic_cast<FlatIndexParams *>(cloned.get());
|
|
ASSERT_NE(cloned_flat, nullptr);
|
|
EXPECT_TRUE(cloned_flat->quantizer_param().enable_rotate());
|
|
|
|
FlatIndexParams params2(MetricType::L2, QuantizeType::INT8);
|
|
EXPECT_FALSE(params == params2);
|
|
}
|
|
|
|
// IVFIndexParams
|
|
{
|
|
IVFIndexParams params(MetricType::IP, 128, 10, false, QuantizeType::INT8);
|
|
EXPECT_FALSE(params.quantizer_param().enable_rotate());
|
|
|
|
params.set_quantizer_param(QuantizerParam(true));
|
|
EXPECT_TRUE(params.quantizer_param().enable_rotate());
|
|
EXPECT_TRUE(params.enable_rotate());
|
|
|
|
auto cloned = params.clone();
|
|
auto *cloned_ivf = dynamic_cast<IVFIndexParams *>(cloned.get());
|
|
ASSERT_NE(cloned_ivf, nullptr);
|
|
EXPECT_TRUE(cloned_ivf->quantizer_param().enable_rotate());
|
|
|
|
IVFIndexParams params2(MetricType::IP, 128, 10, false, QuantizeType::INT8);
|
|
EXPECT_FALSE(params == params2);
|
|
}
|
|
} |