chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:47:42 +08:00
commit be3ef883e1
1214 changed files with 431743 additions and 0 deletions
+49
View File
@@ -0,0 +1,49 @@
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
cc_directory(common)
cc_directories(crash_recovery)
cc_directory(sqlengine)
cc_directories(index)
if(APPLE)
set(APPLE_FRAMEWORK_LIBS
-framework CoreFoundation
-framework CoreGraphics
-framework CoreData
-framework CoreText
-framework Security
-framework Foundation
-Wl,-U,_MallocExtension_ReleaseFreeMemory
-Wl,-U,_ProfilerStart
-Wl,-U,_ProfilerStop
-Wl,-U,_RegisterThriftProtocol
)
endif()
file(GLOB ALL_TEST_SRCS *_test.cc)
foreach(CC_SRCS ${ALL_TEST_SRCS})
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
cc_gmock(
NAME ${CC_TARGET} STRICT
LIBS zvec
zvec_proto
core_knn_flat
core_knn_flat_sparse
core_knn_hnsw
core_knn_hnsw_rabitq
core_knn_hnsw_sparse
core_knn_ivf
core_knn_diskann
core_mix_reducer
core_metric
core_utility
core_quantizer
${CMAKE_THREAD_LIBS_INIT}
${CMAKE_DL_LIBS}
SRCS ${CC_SRCS} index/utils/utils.cc
INCS . .. ../../src
LDFLAGS ${APPLE_FRAMEWORK_LIBS}
)
cc_test_suite(zvec ${CC_TARGET})
endforeach()
File diff suppressed because it is too large Load Diff
+32
View File
@@ -0,0 +1,32 @@
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
if(APPLE)
set(APPLE_FRAMEWORK_LIBS
-framework CoreFoundation
-framework CoreGraphics
-framework CoreData
-framework CoreText
-framework Security
-framework Foundation
-Wl,-U,_MallocExtension_ReleaseFreeMemory
-Wl,-U,_ProfilerStart
-Wl,-U,_ProfilerStop
-Wl,-U,_RegisterThriftProtocol
)
endif()
file(GLOB ALL_TEST_SRCS *_test.cc)
foreach(CC_SRCS ${ALL_TEST_SRCS})
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
cc_gmock(
NAME ${CC_TARGET} STRICT
LIBS zvec_common
${CMAKE_THREAD_LIBS_INIT}
${CMAKE_DL_LIBS}
SRCS ${CC_SRCS}
INCS .. ../../src
LDFLAGS ${APPLE_FRAMEWORK_LIBS}
)
cc_test_suite(zvec_common ${CC_TARGET})
endforeach()
+246
View File
@@ -0,0 +1,246 @@
// 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/config.h"
#include <gtest/gtest.h>
#include "zvec/db/status.h"
using namespace zvec;
class ConfigTest : public ::testing::Test {
protected:
void SetUp() override {
// Reset GlobalConfig for each test
// Note: Since GlobalConfig is a singleton and uses atomic flag,
// we cannot easily reset it. In a real test environment, you might
// need to use a testing framework that supports fixture reset or
// modify the GlobalConfig to support reset for testing purposes.
}
};
TEST_F(ConfigTest, InitializeWithDefaultConfig) {
GlobalConfig::ConfigData config;
// Test initialization with default config
auto status = GlobalConfig::Instance().Initialize(config);
ASSERT_TRUE(status.ok()) << "Initialization failed: " << status.message();
// Verify default values
ASSERT_GT(GlobalConfig::Instance().memory_limit_bytes(), 0);
ASSERT_EQ(GlobalConfig::Instance().log_level(),
GlobalConfig::LogLevel::kWarn);
ASSERT_EQ(GlobalConfig::Instance().log_type(), "ConsoleLogger");
ASSERT_GT(GlobalConfig::Instance().query_thread_count(), 0);
ASSERT_EQ(GlobalConfig::Instance().invert_to_forward_scan_ratio(), 0.9f);
ASSERT_EQ(GlobalConfig::Instance().brute_force_by_keys_ratio(), 0.1f);
ASSERT_EQ(GlobalConfig::Instance().fts_brute_force_by_keys_ratio(), 0.05f);
ASSERT_GT(GlobalConfig::Instance().optimize_thread_count(), 0);
}
TEST_F(ConfigTest, InitializeWithCustomConsoleLogConfig) {
GlobalConfig::ConfigData config;
config.log_config = std::make_shared<GlobalConfig::ConsoleLogConfig>(
GlobalConfig::LogLevel::kDebug);
config.memory_limit_bytes = 1024 * 1024 * 1024; // 1GB
config.query_thread_count = 4;
config.optimize_thread_count = 2;
auto status = GlobalConfig::Instance().Initialize(config);
// First initialization should succeed
if (status.code() == StatusCode::INVALID_ARGUMENT &&
status.message().find("already initialized") != std::string::npos) {
// If already initialized, skip this test
GTEST_SKIP() << "GlobalConfig already initialized";
}
}
TEST_F(ConfigTest, InitializeWithCustomFileLogConfig) {
GlobalConfig::ConfigData config;
auto file_config = std::make_shared<GlobalConfig::FileLogConfig>(
GlobalConfig::LogLevel::kInfo, "/tmp/logs", "test.log", 1024, 14);
config.log_config = file_config;
config.memory_limit_bytes = 2 * 1024 * 1024 * 1024ULL; // 2GB
config.query_thread_count = 8;
config.optimize_thread_count = 4;
auto status = GlobalConfig::Instance().Initialize(config);
// First initialization should succeed
if (status.code() == StatusCode::INVALID_ARGUMENT &&
status.message().find("already initialized") != std::string::npos) {
// If already initialized, skip this test
GTEST_SKIP() << "GlobalConfig already initialized";
}
}
TEST_F(ConfigTest, DoubleInitializationSilentlyFails) {
GlobalConfig::ConfigData config;
auto status1 = GlobalConfig::Instance().Initialize(config);
// If first initialization failed due to already being initialized
if (status1.code() == StatusCode::INVALID_ARGUMENT &&
status1.message().find("already initialized") != std::string::npos) {
// Try again with a fresh config
auto status2 = GlobalConfig::Instance().Initialize(config);
ASSERT_FALSE(status2.ok());
ASSERT_EQ(status2.code(), StatusCode::INVALID_ARGUMENT);
ASSERT_NE(status2.message().find("already initialized"), std::string::npos);
} else {
// First initialization succeeded, second should fail
ASSERT_TRUE(status1.ok());
// The second initialization is allowed but becomes a no-op
auto status2 = GlobalConfig::Instance().Initialize(config);
ASSERT_TRUE(status2.ok());
}
}
TEST_F(ConfigTest, ValidateConfigWithInvalidMemoryLimit) {
GlobalConfig::ConfigData config;
config.memory_limit_bytes = 0; // Invalid value
GlobalConfig
config_instance; // Create a local instance for testing validation
auto status = config_instance.Validate(config);
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
ASSERT_NE(status.message().find("memory_limit_bytes must be greater than"),
std::string::npos);
}
TEST_F(ConfigTest, ValidateConfigWithInvalidQueryThreadCount) {
GlobalConfig::ConfigData config;
config.query_thread_count = 0; // Invalid value
GlobalConfig config_instance;
auto status = config_instance.Validate(config);
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
ASSERT_NE(status.message().find("query_thread_count must be greater than 0"),
std::string::npos);
}
TEST_F(ConfigTest, ValidateConfigWithInvalidRatios) {
GlobalConfig::ConfigData config;
// Test invalid invert_to_forward_scan_ratio
config.invert_to_forward_scan_ratio = -0.1f;
GlobalConfig config_instance;
auto status = config_instance.Validate(config);
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
ASSERT_NE(status.message().find(
"invert_to_forward_scan_ratio must be between 0 and 1"),
std::string::npos);
// Test invalid brute_force_by_keys_ratio
config.invert_to_forward_scan_ratio = 0.9f; // Reset to valid value
config.brute_force_by_keys_ratio = 1.5f; // Invalid value
status = config_instance.Validate(config);
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
ASSERT_NE(status.message().find(
"brute_force_by_keys_ratio must be between 0 and 1"),
std::string::npos);
// Test invalid fts_brute_force_by_keys_ratio
config.brute_force_by_keys_ratio = 0.1f; // Reset to valid value
config.fts_brute_force_by_keys_ratio = -0.5f; // Invalid value
status = config_instance.Validate(config);
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
ASSERT_NE(status.message().find(
"fts_brute_force_by_keys_ratio must be between 0 and 1"),
std::string::npos);
}
TEST_F(ConfigTest, ValidateConfigWithInvalidFileLogSettings) {
GlobalConfig::ConfigData config;
// Test with empty log directory
auto file_config = std::make_shared<GlobalConfig::FileLogConfig>();
file_config->dir = "";
config.log_config = file_config;
GlobalConfig config_instance;
auto status = config_instance.Validate(config);
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
ASSERT_NE(status.message().find("log_dir cannot be empty"),
std::string::npos);
// Test with empty basename
file_config->dir = "/tmp/logs";
file_config->basename = "";
status = config_instance.Validate(config);
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
ASSERT_NE(status.message().find("log_file basename cannot be empty"),
std::string::npos);
// Test with invalid file size
file_config->basename = "test.log";
file_config->file_size = 0;
status = config_instance.Validate(config);
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
ASSERT_NE(status.message().find("log file_size must be greater than"),
std::string::npos);
// Test with invalid overdue days
file_config->file_size = 1024;
file_config->overdue_days = 0;
status = config_instance.Validate(config);
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
ASSERT_NE(status.message().find("log_overdue_days must be greater than 0"),
std::string::npos);
}
TEST_F(ConfigTest, LogLevelEnumValues) {
ASSERT_EQ(static_cast<int>(GlobalConfig::LogLevel::kDebug), 0);
ASSERT_EQ(static_cast<int>(GlobalConfig::LogLevel::kInfo), 1);
ASSERT_EQ(static_cast<int>(GlobalConfig::LogLevel::kWarn), 2);
ASSERT_EQ(static_cast<int>(GlobalConfig::LogLevel::kError), 3);
ASSERT_EQ(static_cast<int>(GlobalConfig::LogLevel::kFatal), 4);
}
TEST_F(ConfigTest, LogConfigPolymorphism) {
auto console_config = std::make_shared<GlobalConfig::ConsoleLogConfig>();
auto file_config = std::make_shared<GlobalConfig::FileLogConfig>();
ASSERT_EQ(console_config->GetLoggerType(), CONSOLE_LOG_TYPE_NAME);
ASSERT_EQ(file_config->GetLoggerType(), FILE_LOG_TYPE_NAME);
}
// jieba_dict_dir is the only ConfigData field that can be written outside
// of Initialize() — language SDKs call set_default_jieba_dict_dir() at
// module-load to register the dict path they bundled. The setter is
// independent of the Initialize() one-shot lifecycle.
TEST_F(ConfigTest, JiebaDictDirSetterIsIndependentOfInitialize) {
auto saved = GlobalConfig::Instance().jieba_dict_dir();
// Setter works regardless of whether Initialize was called.
GlobalConfig::Instance().set_default_jieba_dict_dir("/tmp/zvec/dict-A");
ASSERT_EQ(GlobalConfig::Instance().jieba_dict_dir(), "/tmp/zvec/dict-A");
// Last writer wins.
GlobalConfig::Instance().set_default_jieba_dict_dir("/tmp/zvec/dict-B");
ASSERT_EQ(GlobalConfig::Instance().jieba_dict_dir(), "/tmp/zvec/dict-B");
// Empty clears.
GlobalConfig::Instance().set_default_jieba_dict_dir("");
ASSERT_EQ(GlobalConfig::Instance().jieba_dict_dir(), "");
GlobalConfig::Instance().set_default_jieba_dict_dir(saved);
}
+160
View File
@@ -0,0 +1,160 @@
// 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/status.h"
#include <gtest/gtest.h>
using namespace zvec;
TEST(StatusTest, DefaultConstructor) {
Status status;
EXPECT_TRUE(status.ok());
EXPECT_EQ(status.code(), StatusCode::OK);
EXPECT_EQ(status.message(), "");
}
TEST(StatusTest, ConstructorWithCodeAndMessage) {
std::string msg = "Test error message";
Status status(StatusCode::INVALID_ARGUMENT, msg);
EXPECT_FALSE(status.ok());
EXPECT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
EXPECT_EQ(status.message(), msg);
}
TEST(StatusTest, ConstructorWithRvalueMessage) {
std::string msg = "Test error message";
Status status(StatusCode::NOT_FOUND, std::move(msg));
EXPECT_FALSE(status.ok());
EXPECT_EQ(status.code(), StatusCode::NOT_FOUND);
EXPECT_EQ(status.message(), "Test error message");
}
TEST(StatusTest, CopyConstructor) {
Status original(StatusCode::INTERNAL_ERROR, "Copy test");
Status copy(original);
EXPECT_FALSE(copy.ok());
EXPECT_EQ(copy.code(), StatusCode::INTERNAL_ERROR);
EXPECT_EQ(copy.message(), "Copy test");
EXPECT_EQ(original.code(), copy.code());
EXPECT_EQ(original.message(), copy.message());
}
TEST(StatusTest, CopyAssignment) {
Status original(StatusCode::PERMISSION_DENIED, "Assignment test");
Status assigned;
assigned = original;
EXPECT_FALSE(assigned.ok());
EXPECT_EQ(assigned.code(), StatusCode::PERMISSION_DENIED);
EXPECT_EQ(assigned.message(), "Assignment test");
}
TEST(StatusTest, MoveConstructor) {
Status original(StatusCode::RESOURCE_EXHAUSTED, "Move test");
Status moved(std::move(original));
EXPECT_FALSE(moved.ok());
EXPECT_EQ(moved.code(), StatusCode::RESOURCE_EXHAUSTED);
EXPECT_EQ(moved.message(), "Move test");
}
TEST(StatusTest, MoveAssignment) {
Status original(StatusCode::UNAVAILABLE, "Move assignment test");
Status moved;
moved = std::move(original);
EXPECT_FALSE(moved.ok());
EXPECT_EQ(moved.code(), StatusCode::UNAVAILABLE);
EXPECT_EQ(moved.message(), "Move assignment test");
}
TEST(StatusTest, ComparisonOperators) {
Status status1(StatusCode::INVALID_ARGUMENT, "Error 1");
Status status2(StatusCode::INVALID_ARGUMENT, "Error 1");
Status status3(StatusCode::NOT_FOUND, "Error 2");
Status ok1;
Status ok2;
EXPECT_TRUE(status1 == status2);
EXPECT_FALSE(status1 == status3);
EXPECT_TRUE(ok1 == ok2);
EXPECT_FALSE(status1 == ok1);
EXPECT_FALSE(status1 != status2);
EXPECT_TRUE(status1 != status3);
EXPECT_FALSE(ok1 != ok2);
EXPECT_TRUE(status1 != ok1);
}
TEST(StatusTest, FactoryMethods) {
auto invalid_arg = Status::InvalidArgument("Invalid arg: ", 42);
EXPECT_FALSE(invalid_arg.ok());
EXPECT_EQ(invalid_arg.code(), StatusCode::INVALID_ARGUMENT);
EXPECT_FALSE(invalid_arg.message().empty());
auto not_found = Status::NotFound("Not found: ", "key");
EXPECT_FALSE(not_found.ok());
EXPECT_EQ(not_found.code(), StatusCode::NOT_FOUND);
EXPECT_FALSE(not_found.message().empty());
auto already_exists = Status::AlreadyExists("Already exists: ", "item");
EXPECT_FALSE(already_exists.ok());
EXPECT_EQ(already_exists.code(), StatusCode::ALREADY_EXISTS);
EXPECT_FALSE(already_exists.message().empty());
auto internal_error = Status::InternalError("Internal error: ", "details");
EXPECT_FALSE(internal_error.ok());
EXPECT_EQ(internal_error.code(), StatusCode::INTERNAL_ERROR);
EXPECT_FALSE(internal_error.message().empty());
auto permission_denied =
Status::PermissionDenied("Permission denied for: ", "resource");
EXPECT_FALSE(permission_denied.ok());
EXPECT_EQ(permission_denied.code(), StatusCode::PERMISSION_DENIED);
EXPECT_FALSE(permission_denied.message().empty());
}
TEST(StatusTest, OKFactory) {
auto ok = Status::OK();
EXPECT_TRUE(ok.ok());
EXPECT_EQ(ok.code(), StatusCode::OK);
EXPECT_EQ(ok.message(), "");
}
TEST(StatusTest, CStringConversion) {
Status status(StatusCode::UNKNOWN, "C string test");
EXPECT_STREQ(status.c_str(), "C string test");
Status ok_status;
EXPECT_STREQ(ok_status.c_str(), "");
}
TEST(StatusTest, OutputStreamOperator) {
Status status(StatusCode::INVALID_ARGUMENT, "Stream test");
std::ostringstream oss;
oss << status;
EXPECT_FALSE(oss.str().empty());
EXPECT_NE(oss.str().find(GetDefaultMessage(StatusCode::INVALID_ARGUMENT)),
std::string::npos);
EXPECT_NE(oss.str().find("Stream test"), std::string::npos);
Status ok_status;
std::ostringstream oss2;
oss2 << ok_status;
EXPECT_FALSE(oss2.str().empty());
EXPECT_NE(oss2.str().find("OK"), std::string::npos);
}
+85
View File
@@ -0,0 +1,85 @@
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
# Crash recovery tests rely on fork()+execvp()+kill() to spawn helper processes
# and simulate crashes. These POSIX process APIs are prohibited on iOS.
# On Windows the tests also fail consistently, so skip them there as well.
if(IOS OR WIN32)
return()
endif()
if(APPLE)
set(APPLE_FRAMEWORK_LIBS
-framework CoreFoundation
-framework CoreGraphics
-framework CoreData
-framework CoreText
-framework Security
-framework Foundation
-Wl,-U,_MallocExtension_ReleaseFreeMemory
-Wl,-U,_ProfilerStart
-Wl,-U,_ProfilerStop
-Wl,-U,_RegisterThriftProtocol
)
endif()
# Common libraries
set(CRASH_RECOVERY_COMMON_LIBS
zvec
zvec_proto
core_knn_flat
core_knn_flat_sparse
core_knn_hnsw
core_knn_hnsw_sparse
core_knn_ivf
core_knn_hnsw_rabitq
core_mix_reducer
core_metric
core_utility
core_quantizer
${CMAKE_THREAD_LIBS_INIT}
${CMAKE_DL_LIBS}
)
# Build data_generator executable
cc_binary(
NAME data_generator
LIBS ${CRASH_RECOVERY_COMMON_LIBS}
SRCS data_generator.cc
INCS .. ../../src
LDFLAGS ${APPLE_FRAMEWORK_LIBS}
)
set_target_properties(data_generator PROPERTIES EXCLUDE_FROM_ALL TRUE)
# Build collection_optimizer executable
cc_binary(
NAME collection_optimizer
LIBS ${CRASH_RECOVERY_COMMON_LIBS}
SRCS collection_optimizer.cc
INCS .. ../../src
LDFLAGS ${APPLE_FRAMEWORK_LIBS}
)
set_target_properties(collection_optimizer PROPERTIES EXCLUDE_FROM_ALL TRUE)
# Build test executables
file(GLOB ALL_TEST_SRCS *_test.cc)
foreach(CC_SRCS ${ALL_TEST_SRCS})
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
cc_gmock(
NAME ${CC_TARGET} STRICT
LIBS ${CRASH_RECOVERY_COMMON_LIBS}
SRCS ${CC_SRCS}
INCS .. ../../src
LDFLAGS ${APPLE_FRAMEWORK_LIBS}
)
add_dependencies(${CC_TARGET} data_generator)
add_dependencies(${CC_TARGET} collection_optimizer)
set_tests_properties(${CC_TARGET} PROPERTIES
ENVIRONMENT "TEST_BINARY_DIR=${PROJECT_BINARY_DIR}"
)
cc_test_suite(zvec_crash_recovery ${CC_TARGET})
endforeach()
@@ -0,0 +1,139 @@
// 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.
#ifndef _WIN32
#include <unistd.h>
#endif
#include <filesystem>
#include <zvec/db/collection.h>
#include <zvec/db/options.h>
#include "zvec/ailego/logger/logger.h"
struct Config {
std::string path;
};
bool ParseArgs(int argc, char **argv, Config &config) {
for (int i = 1; i < argc; i++) {
std::string arg = argv[i];
if (arg == "--path" && i + 1 < argc) {
config.path = argv[++i];
} else if (arg == "--help" || arg == "-h") {
return false;
}
}
// Validate required arguments
if (config.path.empty()) {
return false;
}
return true;
}
void PrintUsage(const char *program) {
std::cout << "Usage: " << program << " --path <collection_path>" << std::endl;
std::cout << std::endl;
std::cout << "Arguments:" << std::endl;
std::cout << " --path Path to the collection (required)"
<< std::endl;
}
int main(int argc, char **argv) {
Config config;
// Parse arguments
if (!ParseArgs(argc, argv, config)) {
PrintUsage(argv[0]);
#ifdef __ANDROID__
_exit(1);
#endif
return 1;
}
try {
std::filesystem::path cwd = std::filesystem::current_path();
std::cout << "[collection_optimizer] Current Working Directory: "
<< cwd.string() << std::endl;
} catch (const std::filesystem::filesystem_error &e) {
std::cout << "[collection_optimizer] Failed to get the current working "
"directory: "
<< e.what() << std::endl;
}
std::cout << "Configuration:" << std::endl;
std::cout << " Path: " << config.path << std::endl;
std::cout << std::endl;
// Scope 'result' so its shared_ptr is released before we call _exit().
zvec::Collection::Ptr collection;
{
auto result = zvec::Collection::Open(config.path,
zvec::CollectionOptions{false, true});
if (!result) {
LOG_ERROR("Failed to open collection[%s]: %s", config.path.c_str(),
result.error().c_str());
#ifdef __ANDROID__
_exit(1);
#endif
return -1;
}
collection = result.value();
}
std::cout << "Collection[" << config.path.c_str() << "] opened successfully"
<< std::endl;
// Print initial stats
std::cout << "Initial stats: " << collection->Stats()->to_string_formatted()
<< std::endl;
auto s = collection->Optimize(zvec::OptimizeOptions{2});
if (s.ok()) {
std::cout << "Optimize completed successfully" << std::endl;
// Print final stats
std::cout << "Final stats: " << collection->Stats()->to_string_formatted()
<< std::endl;
#ifdef __ANDROID__
// On Android with c++_static STL, static destructors of glog/gflags
// crash during process teardown. Use _exit() to skip them.
collection->Flush();
collection.reset();
sync();
std::cout.flush();
std::cerr.flush();
fflush(stdout);
fflush(stderr);
_exit(0);
#endif
return 0;
} else {
std::cout << "Optimize failed: " << s.message() << std::endl;
#ifdef __ANDROID__
collection.reset();
sync();
std::cout.flush();
std::cerr.flush();
fflush(stdout);
fflush(stderr);
_exit(1);
#endif
return 1;
}
}
+257
View File
@@ -0,0 +1,257 @@
// 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 <filesystem>
#include <thread>
#ifdef __ANDROID__
#include <unistd.h> // _exit()
#endif
#include <zvec/ailego/internal/platform.h>
#include <zvec/db/collection.h>
#include "zvec/ailego/logger/logger.h"
#include "utility.h"
constexpr int kBatchSize = 20;
constexpr int kBatchDelayMs = 10;
struct Config {
std::string path;
int start_id = 0;
int end_id = 0;
std::string operation; // "insert", "upsert", "update", "delete"
int version = 999999;
};
bool ParseArgs(int argc, char **argv, Config &config) {
for (int i = 1; i < argc; i++) {
std::string arg = argv[i];
if (arg == "--path" && i + 1 < argc) {
config.path = argv[++i];
} else if (arg == "--start" && i + 1 < argc) {
config.start_id = std::stoi(argv[++i]);
} else if (arg == "--end" && i + 1 < argc) {
config.end_id = std::stoi(argv[++i]);
} else if (arg == "--op" && i + 1 < argc) {
config.operation = argv[++i];
} else if (arg == "--version" && i + 1 < argc) {
config.version = std::stoi(argv[++i]);
} else if (arg == "--help" || arg == "-h") {
return false;
}
}
// Validate required arguments
if (config.path.empty() || config.operation.empty() ||
config.start_id >= config.end_id || config.version == 999999) {
return false;
}
// Validate operation
if (config.operation != "insert" && config.operation != "upsert" &&
config.operation != "update" && config.operation != "delete") {
std::cerr << "Error: Invalid operation '" << config.operation
<< "'. Must be 'insert', 'upsert', 'update', or 'delete'."
<< std::endl;
return false;
}
return true;
}
void PrintUsage(const char *program) {
std::cout << "Usage: " << program
<< " --path <collection_path> --start <start_id> --end <end_id> "
"--op <operation>"
<< std::endl;
std::cout << std::endl;
std::cout << "Arguments:" << std::endl;
std::cout << " --path Path to the collection (required)" << std::endl;
std::cout << " --start Starting document ID (inclusive, required)"
<< std::endl;
std::cout << " --end Ending document ID (exclusive, required)"
<< std::endl;
std::cout
<< " --op Operation: insert, upsert, update, or delete (required)"
<< std::endl;
std::cout << " --version Operation: version (required)" << std::endl;
std::cout << std::endl;
std::cout << "Examples:" << std::endl;
std::cout << " # Insert 1000 documents (pk_0 to pk_999)" << std::endl;
std::cout << " " << program
<< " --path ./test_db --start 0 --end 1000 --op insert --version 0"
<< std::endl;
std::cout << std::endl;
std::cout << " # Update documents 1000-1999" << std::endl;
std::cout
<< " " << program
<< " --path ./test_db --start 1000 --end 2000 --op update --version 1"
<< std::endl;
std::cout << std::endl;
std::cout << " # Upsert documents 0-499" << std::endl;
std::cout << " " << program
<< " --path ./test_db --start 0 --end 500 --op upsert --version 2"
<< std::endl;
}
int main(int argc, char **argv) {
Config config;
// Parse arguments
if (!ParseArgs(argc, argv, config)) {
PrintUsage(argv[0]);
#ifdef __ANDROID__
_exit(1);
#endif
return 1;
}
try {
std::filesystem::path cwd = std::filesystem::current_path();
std::cout << "[data_generator] Current Working Directory: " << cwd.string()
<< std::endl;
} catch (const std::filesystem::filesystem_error &e) {
std::cout
<< "[data_generator] Failed to get the current working directory: "
<< e.what() << std::endl;
}
std::cout << "Configuration:" << std::endl;
std::cout << " Path: " << config.path << std::endl;
std::cout << " Range: [" << config.start_id << ", " << config.end_id
<< ")" << std::endl;
std::cout << " Operation: " << config.operation << std::endl;
std::cout << " BatchSize: " << kBatchSize << std::endl;
std::cout << " BatchDelay: " << kBatchDelayMs << "ms" << std::endl;
std::cout << std::endl;
// Scope 'result' so its shared_ptr is released before we call _exit().
// Without scoping, result.value() returns a reference and the copy means
// result still owns a second shared_ptr; collection.reset() would only
// drop the refcount to 1, and _exit() would skip the Collection destructor,
// leaving the WAL dirty.
zvec::Collection::Ptr collection;
{
auto result = zvec::Collection::Open(
config.path, zvec::CollectionOptions{false, true, 4 * 1024 * 1024});
if (!result) {
LOG_ERROR("Failed to open collection[%s]: %s", config.path.c_str(),
result.error().c_str());
#ifdef __ANDROID__
_exit(1);
#endif
return -1;
}
collection = result.value();
}
LOG_INFO("Collection[%s] opened successfully", config.path.c_str());
// Process documents in batches
int total_docs = config.end_id - config.start_id;
int processed = 0;
int batch_num = 0;
int next_progress_threshold = total_docs / 10; // 10% increments
int progress_percent = 0;
while (config.start_id < config.end_id) {
int batch_end = std::min(config.start_id + kBatchSize, config.end_id);
int batch_count = batch_end - config.start_id;
std::vector<zvec::Doc> docs;
docs.reserve(batch_count);
for (int i = config.start_id; i < batch_end; i++) {
docs.push_back(zvec::CreateTestDoc(i, config.version));
}
zvec::Result<zvec::WriteResults> results;
if (config.operation == "insert") {
results = collection->Insert(docs);
} else if (config.operation == "upsert") {
results = collection->Upsert(docs);
} else if (config.operation == "update") {
results = collection->Update(docs);
} else if (config.operation == "delete") {
std::vector<std::string> pks{};
for (const auto &doc : docs) {
pks.emplace_back(doc.pk());
}
results = collection->Delete(pks);
}
if (!results) {
LOG_ERROR("Failed to perform operation[%s], reason: %s",
config.operation.c_str(), results.error().message().c_str());
#ifdef __ANDROID__
collection.reset();
sync();
_exit(1);
#endif
return 1;
}
for (auto &s : results.value()) {
if (!s.ok()) {
LOG_ERROR("Failed to perform operation[%s], reason: %s",
config.operation.c_str(), s.message().c_str());
#ifdef __ANDROID__
collection.reset();
sync();
_exit(1);
#endif
return 1;
}
}
processed += batch_count;
config.start_id = batch_end;
batch_num++;
// Print progress every 10%
if (processed >= next_progress_threshold) {
progress_percent++;
LOG_INFO("Progress: %d (%d/%d documents)", progress_percent * 10,
processed, total_docs);
next_progress_threshold = (progress_percent + 1) * total_docs / 10;
}
// Sleep between batches
if (config.start_id < config.end_id) {
std::this_thread::sleep_for(std::chrono::milliseconds(kBatchDelayMs));
}
}
std::cout << std::endl;
std::cout << "Success! Processed " << processed << " documents in "
<< batch_num << " batches." << std::endl;
#ifdef __ANDROID__
// On Android with c++_static STL, static destructors of glog/gflags
// crash during process teardown. Use _exit() to skip them.
// Flush + close the collection to persist all buffered data (WAL, memtable),
// then sync() to ensure kernel buffers are written to disk before _exit().
collection->Flush();
collection.reset();
sync();
std::cout.flush();
std::cerr.flush();
fflush(stdout);
fflush(stderr);
_exit(0);
#endif
return 0;
}
@@ -0,0 +1,264 @@
// 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 <csignal>
#include <filesystem>
#include <thread>
#include <gtest/gtest.h>
#include <zvec/db/collection.h>
#include <zvec/db/doc.h>
#include <zvec/db/schema.h>
#include "tests/test_util.h"
#include "utility.h"
#ifdef _WIN32
#include <process.h>
#include <windows.h>
typedef HANDLE pid_t;
#define SIGKILL 9
#define WIFEXITED(status) true
#define WEXITSTATUS(status) (status)
#define WIFSIGNALED(status) ((status) != 0)
#endif
namespace zvec {
static std::string optimizer_bin_;
const std::string collection_name_{"optimize_recovery_test"};
const std::string dir_path_{"optimize_recovery_test_db"};
const zvec::CollectionOptions options_{false, true, 256 * 1024};
const int batch_size{50};
const int num_batches{1000};
static std::string LocateOptimizeGenerator() {
return LocateBinary("collection_optimizer");
}
void ExecuteOptimizer(const std::string &path, int kill_after_seconds = -1) {
bool should_crash = kill_after_seconds >= 0;
#ifdef _WIN32
std::string cmd_str = optimizer_bin_ + " --path " + path;
STARTUPINFOA si = {sizeof(si)};
PROCESS_INFORMATION pi;
std::cout << cmd_str << std::endl;
std::vector<char> cmd_buf(cmd_str.begin(), cmd_str.end());
cmd_buf.push_back('\0');
if (!CreateProcessA(NULL, cmd_buf.data(), NULL, NULL, FALSE, 0, NULL, NULL,
&si, &pi)) {
FAIL() << "CreateProcess failed (" << GetLastError() << ")";
}
if (should_crash) {
std::this_thread::sleep_for(std::chrono::seconds(kill_after_seconds));
DWORD wait_result = WaitForSingleObject(pi.hProcess, 0);
if (wait_result == WAIT_TIMEOUT) {
TerminateProcess(pi.hProcess, 1);
}
}
WaitForSingleObject(pi.hProcess, INFINITE);
DWORD exit_code;
GetExitCodeProcess(pi.hProcess, &exit_code);
CloseHandle(pi.hProcess);
CloseHandle(pi.hThread);
if (!should_crash) {
ASSERT_EQ(exit_code, 0) << "Process failed with exit code: " << exit_code;
}
#else
pid_t pid = fork();
ASSERT_GE(pid, 0);
if (pid == 0) {
char arg_path[] = "--path";
char *args[] = {const_cast<char *>(optimizer_bin_.c_str()), arg_path,
const_cast<char *>(path.c_str()), nullptr};
execvp(args[0], args);
perror("execvp failed");
_exit(1);
}
int status;
if (should_crash) {
std::this_thread::sleep_for(std::chrono::seconds(kill_after_seconds));
if (kill(pid, 0) == 0) {
kill(pid, SIGKILL);
}
waitpid(pid, &status, 0);
} else {
waitpid(pid, &status, 0);
ASSERT_TRUE(WIFEXITED(status))
<< "Child process did not exit normally. Terminated by signal?";
int exit_code = WEXITSTATUS(status);
ASSERT_EQ(exit_code, 0) << "optimizer failed with exit code: " << exit_code;
}
#endif
}
void RunOptimizer(const std::string &path) {
ExecuteOptimizer(path);
}
void RunOptimizerAndCrash(const std::string &path, int seconds) {
ExecuteOptimizer(path, seconds);
}
class OptimizeRecoveryTest : public ::testing::Test {
protected:
void SetUp() override {
zvec::test_util::RemoveTestPath("./optimize_recovery_test_db");
ASSERT_NO_THROW(optimizer_bin_ = LocateOptimizeGenerator());
}
void TearDown() override {
zvec::test_util::RemoveTestPath("./optimize_recovery_test_db");
}
};
TEST_F(OptimizeRecoveryTest, CrashDuringOptimize) {
{ // Create a collection and insert some documents
auto schema = CreateTestSchema(collection_name_);
auto result = Collection::CreateAndOpen(dir_path_, *schema, options_);
ASSERT_TRUE(result.has_value());
auto collection = result.value();
for (int batch = 0; batch < num_batches; batch++) {
std::vector<Doc> docs;
for (int i = 0; i < batch_size; i++) {
docs.push_back(CreateTestDoc(batch * batch_size + i, 0));
}
auto write_result = collection->Insert(docs);
ASSERT_TRUE(write_result);
for (auto &s : write_result.value()) {
ASSERT_TRUE(s.ok());
}
}
ASSERT_EQ(collection->Stats()->doc_count, num_batches * batch_size);
collection.reset();
}
RunOptimizerAndCrash(dir_path_, 4);
{ // Open the collection and verify data integrity
auto result = Collection::Open(dir_path_, options_);
ASSERT_TRUE(result.has_value())
<< "Failed to reopen collection after crash. "
"Recovery mechanism may be broken.";
auto collection = result.value();
uint64_t doc_count{collection->Stats().value().doc_count};
ASSERT_EQ(doc_count, num_batches * batch_size);
for (uint64_t doc_id = 0; doc_id < doc_count; doc_id++) {
Doc expected_doc = CreateTestDoc(doc_id, 0);
std::vector<std::string> pks{};
pks.emplace_back(expected_doc.pk());
if (auto res = collection->Fetch(pks); res) {
auto map = res.value();
if (map.find(expected_doc.pk()) == map.end()) {
FAIL() << "Returned map does not contain doc[" << expected_doc.pk()
<< "]";
}
const auto actual_doc = map.at(expected_doc.pk());
ASSERT_EQ(*actual_doc, expected_doc)
<< "Data mismatch for doc[" << expected_doc.pk() << "]";
} else {
FAIL() << "Failed to fetch doc[" << expected_doc.pk() << "]";
}
}
SearchQuery query;
query.topk_ = 10;
std::vector<float> feature(128, 0.0);
query.target_.set_vector(std::string((const char *)feature.data(),
feature.size() * sizeof(float)));
query.target_.field_name_ = "dense_fp32_field";
auto query_result = collection->Query(query);
ASSERT_TRUE(query_result);
auto doc_list = query_result.value();
ASSERT_EQ(doc_list.size(), 10);
ASSERT_EQ(doc_list[0]->pk(), "pk_0");
// Insert some more documents
for (int batch = num_batches; batch < num_batches + 500; batch++) {
std::vector<Doc> docs;
for (int i = 0; i < batch_size; i++) {
docs.push_back(CreateTestDoc(batch * batch_size + i, 0));
}
auto write_result = collection->Insert(docs);
ASSERT_TRUE(write_result);
for (auto &s : write_result.value()) {
ASSERT_TRUE(s.ok());
}
}
collection.reset();
}
RunOptimizer(dir_path_);
// Open the collection and verify data integrity
auto result = Collection::Open(dir_path_, options_);
ASSERT_TRUE(result.has_value()) << "Failed to reopen collection after crash. "
"Recovery mechanism may be broken.";
auto collection = result.value();
uint64_t doc_count{collection->Stats().value().doc_count};
ASSERT_EQ(doc_count, (num_batches + 500) * batch_size);
for (uint64_t doc_id = 0; doc_id < doc_count; doc_id++) {
Doc expected_doc = CreateTestDoc(doc_id, 0);
std::vector<std::string> pks{};
pks.emplace_back(expected_doc.pk());
if (auto res = collection->Fetch(pks); res) {
auto map = res.value();
if (map.find(expected_doc.pk()) == map.end()) {
FAIL() << "Returned map does not contain doc[" << expected_doc.pk()
<< "]";
}
const auto actual_doc = map.at(expected_doc.pk());
ASSERT_EQ(*actual_doc, expected_doc)
<< "Data mismatch for doc[" << expected_doc.pk() << "]";
} else {
FAIL() << "Failed to fetch doc[" << expected_doc.pk() << "]";
}
}
SearchQuery query;
query.topk_ = 10;
std::vector<float> feature(128, 0.0);
query.target_.set_vector(std::string((const char *)feature.data(),
feature.size() * sizeof(float)));
query.target_.field_name_ = "dense_fp32_field";
auto query_result = collection->Query(query);
ASSERT_TRUE(query_result);
auto doc_list = query_result.value();
ASSERT_EQ(doc_list.size(), 10);
ASSERT_EQ(doc_list[0]->pk(), "pk_0");
}
} // namespace zvec
+205
View File
@@ -0,0 +1,205 @@
// 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.
#pragma once
#include <cstdlib>
#include <filesystem>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>
#include <zvec/db/collection.h>
#include <zvec/db/doc.h>
namespace zvec {
/**
* @brief Create a test schema with deterministic field definitions.
*
* @param name The collection name (default: "crash_recovery_test")
* @return CollectionSchema::Ptr The test schema
*/
inline CollectionSchema::Ptr CreateTestSchema(
const std::string &name = "crash_recovery_test") {
auto schema = std::make_shared<CollectionSchema>(name);
schema->set_max_doc_count_per_segment(10000);
schema->add_field(
std::make_shared<FieldSchema>("int32_field", DataType::INT32, false));
schema->add_field(
std::make_shared<FieldSchema>("int64_field", DataType::INT64, true));
schema->add_field(
std::make_shared<FieldSchema>("float_field", DataType::FLOAT, true));
schema->add_field(
std::make_shared<FieldSchema>("string_field", DataType::STRING, false));
schema->add_field(
std::make_shared<FieldSchema>("bool_field", DataType::BOOL, false));
schema->add_field(std::make_shared<FieldSchema>("array_int32_field",
DataType::ARRAY_INT32, true));
schema->add_field(std::make_shared<FieldSchema>(
"array_string_field", DataType::ARRAY_STRING, false));
schema->add_field(std::make_shared<FieldSchema>(
"dense_fp32_field", DataType::VECTOR_FP32, 128, false,
std::make_shared<HnswIndexParams>(MetricType::L2)));
schema->add_field(std::make_shared<FieldSchema>(
"sparse_fp32_field", DataType::SPARSE_VECTOR_FP32, 0, false,
std::make_shared<HnswIndexParams>(MetricType::IP)));
return schema;
}
/**
* @brief Create a test document with deterministic values based on doc_id.
*
* Document pattern:
* - pk: "pk_{doc_id}"
* - int32_field: doc_id (cast to int32)
* - int64_field: doc_id, null if doc_id % 60 == 0
* - float_field: doc_id / 1000.0, null if doc_id % 70 == 0
* - string_field: "{version}_{doc_id}"
* - bool_field: doc_id % 2 == 0 or flipped if version % 2 !=0
* - array_int32_field: [doc_id, doc_id+1, doc_id+2], null if doc_id % 100 == 0
* - array_string_field: ["str_{version}_0", ...]
* - dense_fp32_field: vector where dense[i] = (doc_id + i) / 1000.0f
* - sparse_fp32_field: sparse vector with indices [0, 10, ...]
*
* @param doc_id The document ID (determines all field values)
* @param version The version of the document
* @return Doc The created document
*/
inline Doc CreateTestDoc(uint64_t doc_id, int version) {
Doc doc;
// Set primary key
std::string pk = "pk_" + std::to_string(doc_id);
doc.set_pk(pk);
// Set scalar fields
doc.set<int32_t>("int32_field", static_cast<int32_t>(doc_id));
// int64_field: nullable, null if doc_id % 60 == 0
if (doc_id % 60 != 0) {
doc.set<int64_t>("int64_field", static_cast<int64_t>(doc_id));
}
// float_field: nullable, null if doc_id % 70 == 0
if (doc_id % 70 != 0) {
doc.set<float>("float_field", static_cast<float>(doc_id) / 1000.0f);
}
// string_field: "value_{id}" or "updated_value_{id}"
std::string string_value =
std::to_string(version) + "_" + std::to_string(doc_id);
doc.set<std::string>("string_field", string_value);
// bool_field: alternating based on doc_id, flipped if updated
bool bool_value = (doc_id % 2 == 0);
if (version % 2 != 0) {
bool_value = !bool_value;
}
doc.set<bool>("bool_field", bool_value);
// array_int32_field: nullable, null if doc_id % 100 == 0
if (doc_id % 100 != 0) {
std::vector<int32_t> array_int32;
for (int i = 0; i < 3; i++) {
array_int32.push_back(static_cast<int32_t>(doc_id + i));
}
doc.set<std::vector<int32_t>>("array_int32_field", array_int32);
}
// array_string_field: ["str_0", "str_1", ...] or ["updated_str_0", ...]
std::vector<std::string> array_string;
size_t array_size = doc_id % 5 + 1; // 1 to 5 elements
for (size_t i = 0; i < array_size; i++) {
array_string.push_back("str_" + std::to_string(version) + "_" +
std::to_string(i));
}
doc.set<std::vector<std::string>>("array_string_field", array_string);
// dense_fp32_field: deterministic pattern
std::vector<float> dense(128);
for (int i = 0; i < 128; i++) {
dense[i] = static_cast<float>(doc_id + i) / 1000.0f;
}
doc.set<std::vector<float>>("dense_fp32_field", dense);
// sparse_fp32_field: sparse vector with indices [0, 10, 20, ..., 100]
// Values based on doc_id: value = (doc_id + index) / 1000.0
std::vector<uint32_t> sparse_indices;
std::vector<float> sparse_values;
for (uint32_t idx = 0; idx <= 100; idx += 10) {
sparse_indices.push_back(idx);
sparse_values.push_back(static_cast<float>(doc_id + idx) / 1000.0f);
}
doc.set<std::pair<std::vector<uint32_t>, std::vector<float>>>(
"sparse_fp32_field", std::make_pair(sparse_indices, sparse_values));
return doc;
}
/**
* @brief Locate a binary by name, searching common paths and TEST_BINARY_DIR.
*
* @param binary_name The base name of the binary (e.g. "data_generator")
* @return std::string The canonical path to the found binary
* @throws std::runtime_error if the binary is not found
*/
inline std::string LocateBinary(const std::string &binary_name) {
namespace fs = std::filesystem;
std::cout << "Current path: " << fs::current_path() << std::endl;
std::vector<std::string> candidates;
const std::vector<std::string> search_paths = {"./", "./bin/"};
for (const auto &p : search_paths) {
candidates.push_back(p);
}
#ifdef _WIN32
for (const auto &p : search_paths) {
candidates.push_back(p + "Debug/");
candidates.push_back(p + "Release/");
}
#endif
const char *test_binary_dir = std::getenv("TEST_BINARY_DIR");
if (test_binary_dir != nullptr) {
candidates.push_back(std::string(test_binary_dir) + "/");
candidates.push_back(std::string(test_binary_dir) + "/bin/");
}
for (auto &p : candidates) {
p += binary_name;
#ifdef _WIN32
p += ".exe";
#endif
}
for (const auto &p : candidates) {
if (fs::exists(p)) {
return fs::canonical(p).string();
}
}
throw std::runtime_error(binary_name + " binary not found");
}
} // namespace zvec
@@ -0,0 +1,418 @@
// 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 <csignal>
#include <filesystem>
#include <thread>
#include <gtest/gtest.h>
#include <zvec/db/collection.h>
#include <zvec/db/doc.h>
#include <zvec/db/schema.h>
#include "tests/test_util.h"
#include "utility.h"
#ifdef _WIN32
#include <process.h>
#include <windows.h>
typedef HANDLE pid_t;
#define SIGKILL 9
#define WIFEXITED(status) true
#define WEXITSTATUS(status) (status)
#define WIFSIGNALED(status) ((status) != 0)
#endif
namespace zvec {
static std::string data_generator_bin_;
const std::string collection_name_{"write_recovery_test"};
const std::string dir_path_{"write_recovery_test_db"};
const zvec::CollectionOptions options_{false, true, 256 * 1024};
static std::string LocateDataGenerator() {
return LocateBinary("data_generator");
}
/**
* Internal helper to execute the data generator process.
* If 'kill_after_seconds' is greater than 0, the process will be forcibly
* terminated after the specified duration to simulate a crash.
*/
void ExecuteProcess(const std::string &start, const std::string &end,
const std::string &op, const std::string &version,
int kill_after_seconds = -1) {
bool should_crash = kill_after_seconds >= 0;
#ifdef _WIN32
// 1. Build the command line string with quotes to handle paths with spaces
std::string cmd_str = data_generator_bin_ + " --path " + dir_path_ +
" --start " + start + " --end " + end + " --op " + op +
" --version " + version;
STARTUPINFOA si = {sizeof(si)};
PROCESS_INFORMATION pi;
std::cout << cmd_str << std::endl;
std::vector<char> cmd_buf(cmd_str.begin(), cmd_str.end());
cmd_buf.push_back('\0');
if (!CreateProcessA(NULL, cmd_buf.data(), NULL, NULL, FALSE, 0, NULL, NULL,
&si, &pi)) {
FAIL() << "CreateProcess failed (" << GetLastError() << ")";
}
if (should_crash) {
std::this_thread::sleep_for(std::chrono::seconds(kill_after_seconds));
// Simulate a crash by killing the process
TerminateProcess(pi.hProcess, 1);
}
WaitForSingleObject(pi.hProcess, INFINITE);
DWORD exit_code;
GetExitCodeProcess(pi.hProcess, &exit_code);
CloseHandle(pi.hProcess);
CloseHandle(pi.hThread);
if (should_crash) {
// In crash tests, we expect the process to have been terminated
ASSERT_NE(exit_code, 0) << "Process was expected to crash/terminate.";
} else {
ASSERT_EQ(exit_code, 0) << "Process failed with exit code: " << exit_code;
}
#else
// POSIX Implementation
pid_t pid = fork();
ASSERT_GE(pid, 0);
if (pid == 0) { // Child process
// Use a vector to manage arguments cleanly
std::vector<const char *> args = {data_generator_bin_.c_str(),
"--path",
dir_path_.c_str(),
"--start",
start.c_str(),
"--end",
end.c_str(),
"--op",
op.c_str(),
"--version",
version.c_str(),
nullptr};
execvp(args[0], const_cast<char *const *>(args.data()));
perror("execvp failed");
_exit(1);
}
int status;
if (should_crash) {
std::this_thread::sleep_for(std::chrono::seconds(kill_after_seconds));
kill(pid, SIGKILL);
waitpid(pid, &status, 0);
ASSERT_TRUE(WIFSIGNALED(status)) << "Process did not crash as expected.";
} else {
waitpid(pid, &status, 0);
ASSERT_TRUE(WIFEXITED(status)) << "Process exited abnormally.";
ASSERT_EQ(WEXITSTATUS(status), 0);
}
#endif
}
// Public API 1: Normal execution
void RunGenerator(const std::string &start, const std::string &end,
const std::string &op, const std::string &version) {
ExecuteProcess(start, end, op, version);
}
// Public API 2: Execution with simulated crash
void RunGeneratorAndCrash(const std::string &start, const std::string &end,
const std::string &op, const std::string &version,
int seconds) {
ExecuteProcess(start, end, op, version, seconds);
}
class CrashRecoveryTest : public ::testing::Test {
protected:
void SetUp() override {
zvec::test_util::RemoveTestPath("./write_recovery_test_db");
ASSERT_NO_THROW(data_generator_bin_ = LocateDataGenerator());
}
void TearDown() override {
zvec::test_util::RemoveTestPath("./write_recovery_test_db");
}
};
TEST_F(CrashRecoveryTest, BasicInsertAndReopen) {
{
auto schema = CreateTestSchema(collection_name_);
auto result = Collection::CreateAndOpen(dir_path_, *schema, options_);
ASSERT_TRUE(result.has_value()) << result.error().message();
auto collection = result.value();
collection.reset();
}
RunGenerator("0", "5000", "insert", "0");
auto result = Collection::Open(dir_path_, options_);
ASSERT_TRUE(result.has_value());
auto collection = result.value();
ASSERT_EQ(collection->Stats().value().doc_count, 5000)
<< "Document count mismatch";
}
TEST_F(CrashRecoveryTest, CrashRecoveryDuringInsertion) {
{
auto schema = CreateTestSchema(collection_name_);
auto result = Collection::CreateAndOpen(dir_path_, *schema, options_);
ASSERT_TRUE(result.has_value()) << result.error().message();
auto collection = result.value();
collection.reset();
}
RunGeneratorAndCrash("0", "10000", "insert", "0", 3);
auto result = Collection::Open(dir_path_, options_);
ASSERT_TRUE(result.has_value()) << "Failed to reopen collection after crash. "
"Recovery mechanism may be broken.";
auto collection = result.value();
uint64_t doc_count{collection->Stats().value().doc_count};
ASSERT_GT(doc_count, 800)
<< "Document count is too low after 3s of insertion and recovery";
for (uint64_t doc_id = 0; doc_id < doc_count; doc_id++) {
const auto expected_doc = CreateTestDoc(doc_id, 0);
std::vector<std::string> pks{};
pks.emplace_back(expected_doc.pk());
if (auto res = collection->Fetch(pks); res) {
auto map = res.value();
if (map.find(expected_doc.pk()) == map.end()) {
FAIL() << "Returned map does not contain doc[" << expected_doc.pk()
<< "]";
}
const auto actual_doc = map.at(expected_doc.pk());
ASSERT_EQ(*actual_doc, expected_doc)
<< "Data mismatch for doc[" << expected_doc.pk() << "]";
} else {
FAIL() << "Failed to fetch doc[" << expected_doc.pk() << "]";
}
}
}
TEST_F(CrashRecoveryTest, CrashRecoveryDuringUpsert) {
{
auto schema = CreateTestSchema(collection_name_);
auto result = Collection::CreateAndOpen(dir_path_, *schema, options_);
ASSERT_TRUE(result.has_value()) << result.error().message();
auto collection = result.value();
collection.reset();
}
RunGenerator("0", "5000", "insert", "0");
{
auto result = Collection::Open(dir_path_, options_);
ASSERT_TRUE(result.has_value()) << result.error().message();
auto collection = result.value();
ASSERT_EQ(collection->Stats().value().doc_count, 5000)
<< "Document count mismatch";
}
RunGeneratorAndCrash("4500", "20000", "upsert", "1", 5);
auto result = Collection::Open(dir_path_, options_);
ASSERT_TRUE(result.has_value()) << "Failed to reopen collection after crash. "
"Recovery mechanism may be broken.";
auto collection = result.value();
uint64_t doc_count{collection->Stats().value().doc_count};
ASSERT_GT(doc_count, 6000)
<< "Document count is too low after 5s of insertion and recovery";
for (uint64_t doc_id = 0; doc_id < doc_count; doc_id++) {
Doc expected_doc;
if (doc_id < 4500) {
expected_doc = CreateTestDoc(doc_id, 0);
} else {
expected_doc = CreateTestDoc(doc_id, 1);
}
std::vector<std::string> pks{};
pks.emplace_back(expected_doc.pk());
if (auto res = collection->Fetch(pks); res) {
auto map = res.value();
if (map.find(expected_doc.pk()) == map.end()) {
FAIL() << "Returned map does not contain doc[" << expected_doc.pk()
<< "]";
}
const auto actual_doc = map.at(expected_doc.pk());
ASSERT_EQ(*actual_doc, expected_doc)
<< "Data mismatch for doc[" << expected_doc.pk() << "]";
} else {
FAIL() << "Failed to fetch doc[" << expected_doc.pk() << "]";
}
}
}
TEST_F(CrashRecoveryTest, CrashRecoveryDuringUpdate) {
{
auto schema = CreateTestSchema(collection_name_);
auto result = Collection::CreateAndOpen(dir_path_, *schema, options_);
ASSERT_TRUE(result.has_value()) << result.error().message();
auto collection = result.value();
collection.reset();
}
RunGenerator("0", "18000", "upsert", "0");
{
auto result = Collection::Open(dir_path_, options_);
ASSERT_TRUE(result.has_value()) << result.error().message();
auto collection = result.value();
ASSERT_EQ(collection->Stats().value().doc_count, 18000)
<< "Document count mismatch";
}
RunGeneratorAndCrash("3000", "15000", "update", "3", 4);
auto result = Collection::Open(dir_path_, options_);
ASSERT_TRUE(result.has_value()) << "Failed to reopen collection after crash. "
"Recovery mechanism may be broken.";
auto collection = result.value();
uint64_t doc_count{collection->Stats().value().doc_count};
ASSERT_EQ(doc_count, 18000) << "Document count mismatch after crash recovery";
// Verify docs before the update range are untouched
for (int doc_id = 0; doc_id < 3000; doc_id++) {
Doc expected_doc = CreateTestDoc(doc_id, 0);
std::vector<std::string> pks{expected_doc.pk()};
auto res = collection->Fetch(pks);
ASSERT_TRUE(res.has_value())
<< "Failed to fetch doc[" << expected_doc.pk() << "]";
auto map = res.value();
ASSERT_NE(map.find(expected_doc.pk()), map.end())
<< "Returned map does not contain doc[" << expected_doc.pk() << "]";
const auto actual_doc = map.at(expected_doc.pk());
ASSERT_EQ(*actual_doc, expected_doc)
<< "Data mismatch for doc[" << expected_doc.pk() << "]";
}
// For docs in the update range [3000, 15000), verify consistency:
// Updates are sequential, so there must be a boundary where all docs
// before it are version 3 (updated) and all docs after are version 0.
bool found_boundary = false;
int boundary_id = -1;
for (int doc_id = 3000; doc_id < 15000; doc_id++) {
std::string pk = "pk_" + std::to_string(doc_id);
std::vector<std::string> pks{pk};
auto res = collection->Fetch(pks);
ASSERT_TRUE(res.has_value()) << "Failed to fetch doc[" << pk << "]";
auto map = res.value();
ASSERT_NE(map.find(pk), map.end())
<< "Returned map does not contain doc[" << pk << "]";
const auto actual_doc = map.at(pk);
Doc version0_doc = CreateTestDoc(doc_id, 0);
Doc version3_doc = CreateTestDoc(doc_id, 3);
if (!found_boundary) {
if (*actual_doc == version0_doc) {
// Found the boundary: this doc was NOT updated
found_boundary = true;
boundary_id = doc_id;
} else {
ASSERT_EQ(*actual_doc, version3_doc)
<< "Doc[" << pk << "] is neither version 0 nor version 3";
}
} else {
// After boundary, all docs must be version 0 (not yet updated)
ASSERT_EQ(*actual_doc, version0_doc)
<< "Doc[" << pk << "] after boundary (" << boundary_id
<< ") should be version 0 but isn't";
}
}
// The crash happened after 4s; at least some docs should have been updated
if (!found_boundary) {
// All docs in range were updated (crash happened very late) - that's fine
} else {
ASSERT_GT(boundary_id, 3000)
<< "No documents were updated before crash - process may not have "
"started in time";
}
}
TEST_F(CrashRecoveryTest, CrashRecoveryDuringDelete) {
{
auto schema = CreateTestSchema(collection_name_);
auto result = Collection::CreateAndOpen(dir_path_, *schema, options_);
ASSERT_TRUE(result.has_value()) << result.error().message();
auto collection = result.value();
collection.reset();
}
RunGenerator("0", "18000", "insert", "0");
{
auto result = Collection::Open(dir_path_, options_);
ASSERT_TRUE(result.has_value()) << result.error().message();
auto collection = result.value();
ASSERT_EQ(collection->Stats().value().doc_count, 18000)
<< "Document count mismatch";
}
RunGeneratorAndCrash("3000", "15000", "delete", "0", 4);
auto result = Collection::Open(dir_path_, options_);
ASSERT_TRUE(result.has_value()) << "Failed to reopen collection after crash. "
"Recovery mechanism may be broken.";
auto collection = result.value();
uint64_t doc_count{collection->Stats().value().doc_count};
ASSERT_LT(doc_count, 18000)
<< "No deletes appear to have been applied before the crash";
ASSERT_GT(doc_count, 6000)
<< "Too many documents deleted, recovery likely lost data";
for (int doc_id = 0; doc_id < 3500; doc_id++) {
auto expected_doc = CreateTestDoc(doc_id, 0);
std::vector<std::string> pks{};
pks.emplace_back(expected_doc.pk());
if (auto res = collection->Fetch(pks); res) {
auto map = res.value();
auto it = map.find(expected_doc.pk());
ASSERT_NE(it, map.end())
<< "Fetch result missing requested pk[" << expected_doc.pk() << "]";
if (doc_id < 3000) {
ASSERT_NE(it->second, nullptr)
<< "Existing doc returned as nullptr [" << expected_doc.pk() << "]";
const auto actual_doc = map.at(expected_doc.pk());
ASSERT_EQ(*actual_doc, expected_doc)
<< "Data mismatch for doc[" << expected_doc.pk() << "]";
} else {
ASSERT_EQ(it->second, nullptr)
<< "Returned doc for deleted pk[" << expected_doc.pk() << "]";
}
} else {
FAIL() << "Failed to fetch doc[" << expected_doc.pk() << "]";
}
}
}
} // namespace zvec
+230
View File
@@ -0,0 +1,230 @@
// 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 <memory>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include <zvec/ailego/io/file.h>
#include <zvec/ailego/utility/file_helper.h>
#include "db/common/file_helper.h"
#include "zvec/db/collection.h"
#include "zvec/db/doc.h"
#include "zvec/db/index_params.h"
#include "zvec/db/options.h"
#include "zvec/db/schema.h"
#include "zvec/db/status.h"
#include "zvec/db/type.h"
using namespace zvec;
static const std::string kTestPath = "./test_fts_query";
class FtsQueryTest : public ::testing::Test {
protected:
void SetUp() override {
FileHelper::RemoveDirectory(kTestPath);
}
void TearDown() override {
FileHelper::RemoveDirectory(kTestPath);
}
// Create a schema with one STRING field (for forward) and one FTS field.
static CollectionSchema::Ptr CreateFtsSchema() {
auto schema = std::make_shared<CollectionSchema>("fts_demo");
// A simple scalar field for forward store
schema->add_field(std::make_shared<FieldSchema>("title", DataType::STRING));
// FTS indexed field
schema->add_field(
std::make_shared<FieldSchema>("content", DataType::STRING, false,
std::make_shared<FtsIndexParams>()));
// A vector field is required for Collection to work (segment open expects
// at least one vector field in the normal schema path).
schema->add_field(std::make_shared<FieldSchema>(
"vec", DataType::VECTOR_FP32, 4, false,
std::make_shared<FlatIndexParams>(MetricType::IP)));
return schema;
}
static Doc MakeDoc(uint64_t id, const std::string &title,
const std::string &content) {
Doc doc;
doc.set_pk("pk_" + std::to_string(id));
doc.set<std::string>("title", title);
doc.set<std::string>("content", content);
// dummy vector
doc.set<std::vector<float>>("vec", std::vector<float>(4, float(id + 0.1)));
return doc;
}
};
TEST_F(FtsQueryTest, BasicFtsQuery) {
auto schema = CreateFtsSchema();
CollectionOptions options;
options.read_only_ = false;
auto result = Collection::CreateAndOpen(kTestPath, *schema, options);
ASSERT_TRUE(result.has_value()) << result.error().message();
auto col = result.value();
// Insert documents
std::vector<Doc> docs;
docs.push_back(MakeDoc(0, "intro", "hello world from zvec"));
docs.push_back(MakeDoc(1, "guide", "hello foo bar"));
docs.push_back(MakeDoc(2, "faq", "baz qux nothing here"));
docs.push_back(MakeDoc(3, "tips", "hello hello hello world"));
auto insert_res = col->Insert(docs);
ASSERT_TRUE(insert_res.has_value()) << insert_res.error().message();
// FTS query: search for "hello"
SearchQuery vq;
vq.target_.field_name_ = "content";
vq.topk_ = 10;
FtsClause fts;
fts.query_string_ = "hello";
vq.target_.clause_ = fts;
auto query_res = col->Query(vq);
ASSERT_TRUE(query_res.has_value()) << query_res.error().message();
auto &results = query_res.value();
// Documents 0, 1, 3 contain "hello"; document 2 does not.
ASSERT_GE(results.size(), 2u);
ASSERT_LE(results.size(), 3u);
}
TEST_F(FtsQueryTest, FtsQueryEmptyField) {
auto schema = CreateFtsSchema();
CollectionOptions options;
options.read_only_ = false;
auto result = Collection::CreateAndOpen(kTestPath, *schema, options);
ASSERT_TRUE(result.has_value());
auto col = result.value();
SearchQuery vq;
vq.target_.field_name_ = ""; // empty
vq.topk_ = 10;
FtsClause fts;
fts.query_string_ = "hello";
vq.target_.clause_ = fts;
auto query_res = col->Query(vq);
ASSERT_FALSE(query_res.has_value());
}
TEST_F(FtsQueryTest, FtsQueryNoMatch) {
auto schema = CreateFtsSchema();
CollectionOptions options;
options.read_only_ = false;
auto result = Collection::CreateAndOpen(kTestPath, *schema, options);
ASSERT_TRUE(result.has_value());
auto col = result.value();
std::vector<Doc> docs;
docs.push_back(MakeDoc(0, "intro", "hello world"));
auto insert_res = col->Insert(docs);
ASSERT_TRUE(insert_res.has_value());
SearchQuery vq;
vq.target_.field_name_ = "content";
vq.topk_ = 10;
FtsClause fts;
fts.query_string_ = "nonexistent_term_xyz";
vq.target_.clause_ = fts;
auto query_res = col->Query(vq);
ASSERT_TRUE(query_res.has_value());
ASSERT_EQ(query_res.value().size(), 0u);
}
// Verify that FTS fields do NOT support add/alter/drop column operations.
// The schema change validation only allows basic numeric types [INT32..DOUBLE].
TEST_F(FtsQueryTest, FtsFieldUnsupportedAddColumn) {
auto schema = CreateFtsSchema();
CollectionOptions options;
options.read_only_ = false;
auto result = Collection::CreateAndOpen(kTestPath, *schema, options);
ASSERT_TRUE(result.has_value());
auto col = result.value();
// Insert a document so the collection is non-empty
std::vector<Doc> docs;
docs.push_back(MakeDoc(0, "intro", "hello world"));
auto insert_res = col->Insert(docs);
ASSERT_TRUE(insert_res.has_value());
ASSERT_TRUE(col->Flush().ok());
// Attempt to add a new FTS column — should fail
auto fts_field = std::make_shared<FieldSchema>(
"new_fts", DataType::STRING, true, std::make_shared<FtsIndexParams>());
auto status = col->AddColumn(fts_field, "", AddColumnOptions());
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
}
TEST_F(FtsQueryTest, FtsFieldUnsupportedDropColumn) {
auto schema = CreateFtsSchema();
CollectionOptions options;
options.read_only_ = false;
auto result = Collection::CreateAndOpen(kTestPath, *schema, options);
ASSERT_TRUE(result.has_value());
auto col = result.value();
// Insert a document so the collection is non-empty
std::vector<Doc> docs;
docs.push_back(MakeDoc(0, "intro", "hello world"));
auto insert_res = col->Insert(docs);
ASSERT_TRUE(insert_res.has_value());
ASSERT_TRUE(col->Flush().ok());
// Attempt to drop an existing FTS column — should fail
auto status = col->DropColumn("content");
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
}
TEST_F(FtsQueryTest, FtsFieldUnsupportedAlterColumn) {
auto schema = CreateFtsSchema();
CollectionOptions options;
options.read_only_ = false;
auto result = Collection::CreateAndOpen(kTestPath, *schema, options);
ASSERT_TRUE(result.has_value());
auto col = result.value();
// Insert a document so the collection is non-empty
std::vector<Doc> docs;
docs.push_back(MakeDoc(0, "intro", "hello world"));
auto insert_res = col->Insert(docs);
ASSERT_TRUE(insert_res.has_value());
ASSERT_TRUE(col->Flush().ok());
// Attempt to alter (rename) the FTS column — should fail
auto status = col->AlterColumn("content", "content_renamed", nullptr,
AlterColumnOptions());
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
// Attempt to alter the FTS column with a new schema — should also fail
auto new_fts_field = std::make_shared<FieldSchema>(
"content", DataType::STRING, true, std::make_shared<FtsIndexParams>());
status = col->AlterColumn("content", "", new_fts_field, AlterColumnOptions());
ASSERT_FALSE(status.ok());
ASSERT_EQ(status.code(), StatusCode::INVALID_ARGUMENT);
}
+64
View File
@@ -0,0 +1,64 @@
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
if(APPLE)
set(APPLE_FRAMEWORK_LIBS
-framework CoreFoundation
-framework CoreGraphics
-framework CoreData
-framework CoreText
-framework Security
-framework Foundation
-Wl,-U,_MallocExtension_ReleaseFreeMemory
-Wl,-U,_ProfilerStart
-Wl,-U,_ProfilerStop
-Wl,-U,_RegisterThriftProtocol
)
endif()
file(GLOB_RECURSE ALL_TEST_SRCS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} *_test.cc)
if(MSVC)
#TODO(windows): should include these UTs
set(_MSVC_EXCLUDE_TESTS
"storage/mmap_store_test.cc"
"storage/mem_store_test.cc"
"column/inverted_column/inverted_indexer_util_test.cc"
"segment/segment_test.cc"
)
foreach(_excl ${_MSVC_EXCLUDE_TESTS})
list(FILTER ALL_TEST_SRCS EXCLUDE REGEX "${_excl}")
endforeach()
endif()
foreach(CC_SRCS ${ALL_TEST_SRCS})
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
cc_gmock(
NAME ${CC_TARGET} STRICT
LIBS zvec
zvec_ailego
zvec_proto
core_knn_diskann
core_metric_static
core_utility_static
core_quantizer_static
core_knn_hnsw core_knn_hnsw_sparse sparsehash
core_knn_flat core_knn_flat_sparse core_knn_ivf
core_knn_hnsw_rabitq core_mix_reducer
Arrow::arrow_dataset
${CMAKE_THREAD_LIBS_INIT}
${CMAKE_DL_LIBS}
SRCS ${CC_SRCS} utils/utils.cc
INCS . .. ../../src
LDFLAGS ${APPLE_FRAMEWORK_LIBS}
)
cc_test_suite(zvec_index ${CC_TARGET})
endforeach()
# Inject TEST_SOURCE_DIR for fts_column_indexer_test so it can locate testdata/
if(TARGET fts_column_indexer_test)
target_compile_definitions(fts_column_indexer_test PRIVATE
TEST_SOURCE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/column/fts_column"
JIEBA_DICT_DIR="${PROJECT_SOURCE_DIR}/thirdparty/cppjieba/cppjieba-5.6.7/dict")
endif()
@@ -0,0 +1,265 @@
// 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 "db/index/column/fts_column/tokenizer/ascii_folding_token_filter.h"
#include <string>
#include <vector>
#include <gtest/gtest.h>
using namespace zvec::fts;
static std::vector<Token> make_tokens(const std::vector<std::string> &texts) {
std::vector<Token> tokens;
for (size_t i = 0; i < texts.size(); ++i) {
tokens.push_back({texts[i], 0, static_cast<uint32_t>(i)});
}
return tokens;
}
class AsciiFoldingTokenFilterTest : public ::testing::Test {
protected:
AsciiFoldingTokenFilter filter_;
};
// --- Pure ASCII passthrough ---
TEST_F(AsciiFoldingTokenFilterTest, AsciiPassthrough) {
auto result = filter_.filter(make_tokens({"hello", "world", "123"}));
ASSERT_EQ(result.size(), 3u);
EXPECT_EQ(result[0].text, "hello");
EXPECT_EQ(result[1].text, "world");
EXPECT_EQ(result[2].text, "123");
}
TEST_F(AsciiFoldingTokenFilterTest, EmptyToken) {
auto result = filter_.filter(make_tokens({""}));
EXPECT_TRUE(result.empty());
}
TEST_F(AsciiFoldingTokenFilterTest, EmptyList) {
std::vector<Token> tokens;
auto result = filter_.filter(std::move(tokens));
EXPECT_TRUE(result.empty());
}
// --- Latin diacritics (NFKD + STRIPMARK handles these) ---
TEST_F(AsciiFoldingTokenFilterTest, LatinAccentedVowels) {
// àáâãäå → aaaaaa
auto result = filter_.filter(
make_tokens({"\xC3\xA0\xC3\xA1\xC3\xA2\xC3\xA3\xC3\xA4\xC3\xA5"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "aaaaaa");
}
TEST_F(AsciiFoldingTokenFilterTest, LatinAccentedConsonants) {
// ñ → n, ç → c
auto result = filter_.filter(make_tokens({"\xC3\xB1", "\xC3\xA7"}));
ASSERT_EQ(result.size(), 2u);
EXPECT_EQ(result[0].text, "n");
EXPECT_EQ(result[1].text, "c");
}
TEST_F(AsciiFoldingTokenFilterTest, UppercaseAccented) {
// ÜBER → UBER
auto result =
filter_.filter(make_tokens({"\xC3\x9C"
"BER"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "UBER");
}
// --- Supplementary table entries ---
TEST_F(AsciiFoldingTokenFilterTest, SharpS) {
// ß → ss
auto result = filter_.filter(make_tokens({"\xC3\x9F"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "ss");
}
TEST_F(AsciiFoldingTokenFilterTest, OWithStroke) {
// ø → o
auto result = filter_.filter(make_tokens({"\xC3\xB8"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "o");
}
TEST_F(AsciiFoldingTokenFilterTest, DWithStroke) {
// đ → d
auto result = filter_.filter(make_tokens({"\xC4\x91"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "d");
}
TEST_F(AsciiFoldingTokenFilterTest, LWithStroke) {
// ł → l
auto result = filter_.filter(make_tokens({"\xC5\x82"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "l");
}
TEST_F(AsciiFoldingTokenFilterTest, Eth) {
// ð → d
auto result = filter_.filter(make_tokens({"\xC3\xB0"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "d");
}
TEST_F(AsciiFoldingTokenFilterTest, Thorn) {
// þ → th
auto result = filter_.filter(make_tokens({"\xC3\xBE"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "th");
}
TEST_F(AsciiFoldingTokenFilterTest, OeLigature) {
// œ → oe
auto result = filter_.filter(make_tokens({"\xC5\x93"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "oe");
}
TEST_F(AsciiFoldingTokenFilterTest, AeLigature) {
// Æ → AE, æ → ae
auto result = filter_.filter(make_tokens({"\xC3\x86", "\xC3\xA6"}));
ASSERT_EQ(result.size(), 2u);
EXPECT_EQ(result[0].text, "AE");
EXPECT_EQ(result[1].text, "ae");
}
TEST_F(AsciiFoldingTokenFilterTest, CapitalThorn) {
// Þ → TH
auto result = filter_.filter(make_tokens({"\xC3\x9E"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "TH");
}
TEST_F(AsciiFoldingTokenFilterTest, IjLigature) {
// IJ → IJ, ij → ij
auto result = filter_.filter(make_tokens({"\xC4\xB2", "\xC4\xB3"}));
ASSERT_EQ(result.size(), 2u);
EXPECT_EQ(result[0].text, "IJ");
EXPECT_EQ(result[1].text, "ij");
}
TEST_F(AsciiFoldingTokenFilterTest, CapitalSharpS) {
// ẞ (U+1E9E) → SS
auto result = filter_.filter(make_tokens({"\xE1\xBA\x9E"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "SS");
}
// --- Ligatures handled by NFKD ---
TEST_F(AsciiFoldingTokenFilterTest, FiLigature) {
// fi (U+FB01) → fi
auto result = filter_.filter(make_tokens({"\xEF\xAC\x81"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "fi");
}
TEST_F(AsciiFoldingTokenFilterTest, FlLigature) {
// fl (U+FB02) → fl
auto result = filter_.filter(make_tokens({"\xEF\xAC\x82"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "fl");
}
// --- Fullwidth forms ---
TEST_F(AsciiFoldingTokenFilterTest, FullwidthLatinCapital) {
// (U+FF21) → A
auto result = filter_.filter(make_tokens({"\xEF\xBC\xA1"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "A");
}
TEST_F(AsciiFoldingTokenFilterTest, FullwidthDigit) {
// (U+FF10) → 0
auto result = filter_.filter(make_tokens({"\xEF\xBC\x90"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "0");
}
// --- CJK passthrough (no ASCII equivalent) ---
TEST_F(AsciiFoldingTokenFilterTest, CJKPassthrough) {
// 中文 should remain unchanged
auto result = filter_.filter(make_tokens({"\xE4\xB8\xAD\xE6\x96\x87"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "\xE4\xB8\xAD\xE6\x96\x87");
}
TEST_F(AsciiFoldingTokenFilterTest, GreekTonosPassthrough) {
// Greek ά has no ASCII equivalent, so it should remain unchanged.
auto result = filter_.filter(make_tokens({"\xCE\xAC"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "\xCE\xAC");
}
// --- Mixed content ---
TEST_F(AsciiFoldingTokenFilterTest, MixedAsciiAndAccented) {
// café → cafe
auto result = filter_.filter(make_tokens({"caf\xC3\xA9"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "cafe");
}
TEST_F(AsciiFoldingTokenFilterTest, MixedAsciiAndCJK) {
// hello中文 → hello中文 (ASCII kept, CJK kept)
auto result = filter_.filter(make_tokens({"hello\xE4\xB8\xAD\xE6\x96\x87"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "hello\xE4\xB8\xAD\xE6\x96\x87");
}
TEST_F(AsciiFoldingTokenFilterTest, MixedLatinAndGreekTonos) {
// caféά → cafeά (Latin folds, Greek stays original)
auto result = filter_.filter(make_tokens({"caf\xC3\xA9\xCE\xAC"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "cafe\xCE\xAC");
}
TEST_F(AsciiFoldingTokenFilterTest, DecomposedLatinCombiningMarkPassthrough) {
// Align with ES/Lucene asciifolding: combining marks are not folded by
// themselves, so decomposed "cafe + U+0301" remains unchanged.
auto result = filter_.filter(make_tokens({"cafe\xCC\x81"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "cafe\xCC\x81");
}
// --- Preserves offset and position ---
TEST_F(AsciiFoldingTokenFilterTest, PreservesOffsetAndPosition) {
std::vector<Token> tokens = {{"\xC3\xA9", 10, 5}};
auto result = filter_.filter(std::move(tokens));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "e");
EXPECT_EQ(result[0].offset, 10u);
EXPECT_EQ(result[0].position, 5u);
}
// --- Name ---
TEST_F(AsciiFoldingTokenFilterTest, FilterName) {
EXPECT_STREQ(filter_.name(), "ascii_folding");
}
TEST_F(AsciiFoldingTokenFilterTest, StandaloneCombiningMarkPassthrough) {
// U+0301 COMBINING ACUTE ACCENT has no ASCII equivalent on its own.
auto result = filter_.filter(make_tokens({"\xCC\x81"}));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "\xCC\x81");
}
@@ -0,0 +1,866 @@
// 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 "db/index/column/fts_column/posting/bitpacked_posting_list.h"
#include <algorithm>
#include <cstdint>
#include <cstring>
#include <numeric>
#include <random>
#include <vector>
#include <gtest/gtest.h>
#include "db/index/column/fts_column/bm25_scorer.h"
#include "db/index/column/fts_column/posting/bitpacked_simd_scalar.h"
#if defined(__SSE4_1__)
#include "db/index/column/fts_column/posting/bitpacked_simd_sse41.h"
#endif
using namespace zvec::fts;
// ============================================================
// Helper: create a BM25Scorer with reasonable defaults
// ============================================================
static BM25Scorer make_scorer(uint64_t total_docs = 1000,
uint64_t total_tokens = 50000) {
BM25Scorer scorer;
scorer.update_stats(total_docs, total_tokens);
return scorer;
}
// ============================================================
// bits_needed()
// ============================================================
TEST(BitPackedPostingListTest, BitsNeededZero) {
EXPECT_EQ(BitPackedPostingList::bits_needed(0), 0);
}
TEST(BitPackedPostingListTest, BitsNeededOne) {
EXPECT_EQ(BitPackedPostingList::bits_needed(1), 1);
}
TEST(BitPackedPostingListTest, BitsNeededPowerOfTwo) {
EXPECT_EQ(BitPackedPostingList::bits_needed(2), 2);
EXPECT_EQ(BitPackedPostingList::bits_needed(4), 3);
EXPECT_EQ(BitPackedPostingList::bits_needed(8), 4);
EXPECT_EQ(BitPackedPostingList::bits_needed(256), 9);
EXPECT_EQ(BitPackedPostingList::bits_needed(1024), 11);
}
TEST(BitPackedPostingListTest, BitsNeededNonPowerOfTwo) {
EXPECT_EQ(BitPackedPostingList::bits_needed(3), 2);
EXPECT_EQ(BitPackedPostingList::bits_needed(5), 3);
EXPECT_EQ(BitPackedPostingList::bits_needed(7), 3);
EXPECT_EQ(BitPackedPostingList::bits_needed(255), 8);
EXPECT_EQ(BitPackedPostingList::bits_needed(1023), 10);
}
TEST(BitPackedPostingListTest, BitsNeededMaxUint32) {
EXPECT_EQ(BitPackedPostingList::bits_needed(0xFFFFFFFF), 32);
}
// ============================================================
// pack_uint32 / unpack_uint32 round-trip
// ============================================================
class BitPackingTest : public ::testing::TestWithParam<uint8_t> {};
TEST_P(BitPackingTest, PackUnpackRoundTrip128) {
const uint8_t bitwidth = GetParam();
if (bitwidth == 0) return;
const uint32_t count = 128;
const uint32_t mask =
(bitwidth == 32) ? 0xFFFFFFFFu : ((1u << bitwidth) - 1u);
// Generate test values
std::vector<uint32_t> original(count);
for (uint32_t i = 0; i < count; ++i) {
original[i] = (i * 17 + 3) & mask; // deterministic pattern
}
// Pack
const size_t packed_size =
BitPackedPostingList::packed_byte_size(bitwidth, count);
std::vector<uint8_t> packed(packed_size, 0);
BitPackedPostingList::pack_uint32(original.data(), bitwidth, count,
packed.data());
// Unpack
std::vector<uint32_t> decoded(count, 0);
BitPackedPostingList::unpack_uint32(packed.data(), bitwidth, count,
decoded.data());
// Verify
for (uint32_t i = 0; i < count; ++i) {
EXPECT_EQ(decoded[i], original[i])
<< "Mismatch at index " << i << " with bitwidth " << (int)bitwidth;
}
}
TEST_P(BitPackingTest, PackUnpackRoundTripSmall) {
const uint8_t bitwidth = GetParam();
if (bitwidth == 0) return;
// Test with a small count (not a full block)
const uint32_t count = 7;
const uint32_t mask =
(bitwidth == 32) ? 0xFFFFFFFFu : ((1u << bitwidth) - 1u);
std::vector<uint32_t> original(count);
for (uint32_t i = 0; i < count; ++i) {
original[i] = i & mask;
}
const size_t packed_size =
BitPackedPostingList::packed_byte_size(bitwidth, count);
std::vector<uint8_t> packed(packed_size, 0);
BitPackedPostingList::pack_uint32(original.data(), bitwidth, count,
packed.data());
std::vector<uint32_t> decoded(count, 0);
BitPackedPostingList::unpack_uint32(packed.data(), bitwidth, count,
decoded.data());
for (uint32_t i = 0; i < count; ++i) {
EXPECT_EQ(decoded[i], original[i])
<< "Mismatch at index " << i << " with bitwidth " << (int)bitwidth;
}
}
// The scalar full-block packer must round-trip on its own. This exercises the
// portable (non-SIMD) path directly regardless of the host CPU, so the
// architecture used by ARM / no-SSE builds is always covered even when tests
// run on x86.
TEST_P(BitPackingTest, ScalarFullBlockRoundTrip) {
const uint8_t bitwidth = GetParam();
if (bitwidth == 0) return;
const uint32_t count = 128;
const uint32_t mask =
(bitwidth == 32) ? 0xFFFFFFFFu : ((1u << bitwidth) - 1u);
std::vector<uint32_t> values(count);
for (uint32_t i = 0; i < count; ++i) {
values[i] = (i * 2654435761u + 7u) & mask; // pseudo-random, deterministic
}
alignas(16) uint8_t packed[32 * 16];
simd::scalar_pack_uint32_128(values.data(), bitwidth, packed);
alignas(16) uint32_t decoded[count];
simd::scalar_unpack_uint32_128(packed, bitwidth, decoded);
for (uint32_t i = 0; i < count; ++i) {
EXPECT_EQ(decoded[i], values[i])
<< "Mismatch at index " << i << " with bitwidth "
<< static_cast<int>(bitwidth);
}
}
#if defined(__SSE4_1__)
// Cross-architecture portability guard: the scalar full-block packer must
// produce byte-identical output to the SSE4.1 SIMD packer, and each must be
// able to decode the other's output. If the scalar fallback ever diverged from
// the SIMD layout again (the original cross-arch corruption bug), an index
// built on x86 would be silently mis-decoded on ARM/no-SSE and vice versa.
TEST_P(BitPackingTest, ScalarLayoutMatchesSimd) {
const uint8_t bitwidth = GetParam();
if (bitwidth == 0) return;
const uint32_t count = 128;
const uint32_t mask =
(bitwidth == 32) ? 0xFFFFFFFFu : ((1u << bitwidth) - 1u);
std::vector<uint32_t> values(count);
for (uint32_t i = 0; i < count; ++i) {
values[i] = (i * 40503u + 12345u) & mask; // deterministic
}
const size_t packed_size = static_cast<size_t>(bitwidth) * 16;
alignas(16) uint8_t scalar_packed[32 * 16];
alignas(16) uint8_t simd_packed[32 * 16];
simd::scalar_pack_uint32_128(values.data(), bitwidth, scalar_packed);
simd::sse41_pack_uint32_128(values.data(), bitwidth, simd_packed);
EXPECT_EQ(0, std::memcmp(scalar_packed, simd_packed, packed_size))
<< "Scalar and SIMD packed bytes differ for bitwidth "
<< static_cast<int>(bitwidth) << " — on-disk format is not portable";
// SIMD decodes scalar-packed bytes.
alignas(16) uint32_t simd_decoded[count];
simd::sse41_unpack_uint32_128(scalar_packed, bitwidth, simd_decoded);
// Scalar decodes SIMD-packed bytes.
alignas(16) uint32_t scalar_decoded[count];
simd::scalar_unpack_uint32_128(simd_packed, bitwidth, scalar_decoded);
for (uint32_t i = 0; i < count; ++i) {
EXPECT_EQ(simd_decoded[i], values[i]) << "SIMD-decode-of-scalar @" << i;
EXPECT_EQ(scalar_decoded[i], values[i]) << "scalar-decode-of-SIMD @" << i;
}
}
#endif // defined(__SSE4_1__)
// Test all bitwidths from 1 to 32
INSTANTIATE_TEST_SUITE_P(AllBitwidths, BitPackingTest,
::testing::Range(static_cast<uint8_t>(1),
static_cast<uint8_t>(33)));
TEST(BitPackingTest, PackUnpackZeroBitwidth) {
const uint32_t count = 128;
std::vector<uint32_t> original(count, 0);
std::vector<uint32_t> decoded(count, 99);
// bitwidth 0: all values must be 0
BitPackedPostingList::unpack_uint32(nullptr, 0, count, decoded.data());
for (uint32_t i = 0; i < count; ++i) {
EXPECT_EQ(decoded[i], 0u);
}
}
// ============================================================
// Encode / Decode: empty posting list
// ============================================================
TEST(BitPackedPostingListTest, EncodeDecodeEmpty) {
BM25Scorer scorer = make_scorer();
std::string encoded =
BitPackedPostingList::encode(nullptr, nullptr, nullptr, 0, 0, scorer);
EXPECT_TRUE(BitPackedPostingList::is_bitpacked_format(encoded.data(),
encoded.size()));
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
EXPECT_EQ(iter.cost(), 0u);
EXPECT_EQ(iter.next_doc(), BitPackedPostingIterator::NO_MORE_DOCS);
}
// ============================================================
// Encode / Decode: single element
// ============================================================
TEST(BitPackedPostingListTest, EncodeDecodeSingleElement) {
BM25Scorer scorer = make_scorer();
uint32_t doc_ids[] = {42};
uint32_t tfs[] = {3};
uint32_t doc_lens[] = {100};
std::string encoded =
BitPackedPostingList::encode(doc_ids, tfs, doc_lens, 1, 1, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
EXPECT_EQ(iter.cost(), 1u);
EXPECT_EQ(iter.next_doc(), 42u);
EXPECT_EQ(iter.doc_id(), 42u);
EXPECT_EQ(iter.term_freq(), 3u);
EXPECT_EQ(iter.doc_len(), 100u);
EXPECT_EQ(iter.next_doc(), BitPackedPostingIterator::NO_MORE_DOCS);
}
// ============================================================
// Encode / Decode: small list (< 128)
// ============================================================
TEST(BitPackedPostingListTest, EncodeDecodeSmallList) {
BM25Scorer scorer = make_scorer();
const size_t count = 10;
std::vector<uint32_t> doc_ids(count);
std::vector<uint32_t> tfs(count);
std::vector<uint32_t> doc_lens(count);
for (size_t i = 0; i < count; ++i) {
doc_ids[i] = static_cast<uint32_t>(i * 5);
tfs[i] = static_cast<uint32_t>(i + 1);
doc_lens[i] = static_cast<uint32_t>(50 + i * 10);
}
std::string encoded = BitPackedPostingList::encode(
doc_ids.data(), tfs.data(), doc_lens.data(), count, count, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
EXPECT_EQ(iter.cost(), count);
for (size_t i = 0; i < count; ++i) {
uint32_t doc = iter.next_doc();
EXPECT_EQ(doc, doc_ids[i]) << "Mismatch at index " << i;
EXPECT_EQ(iter.term_freq(), tfs[i]) << "TF mismatch at index " << i;
EXPECT_EQ(iter.doc_len(), doc_lens[i]) << "DocLen mismatch at index " << i;
}
EXPECT_EQ(iter.next_doc(), BitPackedPostingIterator::NO_MORE_DOCS);
}
// ============================================================
// Encode / Decode: exactly 128 elements (one full block)
// ============================================================
TEST(BitPackedPostingListTest, EncodeDecodeExact128) {
BM25Scorer scorer = make_scorer();
const size_t count = 128;
std::vector<uint32_t> doc_ids(count);
std::vector<uint32_t> tfs(count);
std::vector<uint32_t> doc_lens(count);
for (size_t i = 0; i < count; ++i) {
doc_ids[i] = static_cast<uint32_t>(i * 3);
tfs[i] = static_cast<uint32_t>((i % 10) + 1);
doc_lens[i] = static_cast<uint32_t>(100 + i);
}
std::string encoded = BitPackedPostingList::encode(
doc_ids.data(), tfs.data(), doc_lens.data(), count, count, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
EXPECT_EQ(iter.cost(), count);
for (size_t i = 0; i < count; ++i) {
uint32_t doc = iter.next_doc();
EXPECT_EQ(doc, doc_ids[i]) << "Mismatch at index " << i;
EXPECT_EQ(iter.term_freq(), tfs[i]) << "TF mismatch at index " << i;
EXPECT_EQ(iter.doc_len(), doc_lens[i]) << "DocLen mismatch at index " << i;
}
EXPECT_EQ(iter.next_doc(), BitPackedPostingIterator::NO_MORE_DOCS);
}
// ============================================================
// Encode / Decode: 129 elements (two blocks, last block has 1 element)
// ============================================================
TEST(BitPackedPostingListTest, EncodeDecodeCrossBlockBoundary) {
BM25Scorer scorer = make_scorer();
const size_t count = 129;
std::vector<uint32_t> doc_ids(count);
std::vector<uint32_t> tfs(count);
std::vector<uint32_t> doc_lens(count);
for (size_t i = 0; i < count; ++i) {
doc_ids[i] = static_cast<uint32_t>(i * 2);
tfs[i] = static_cast<uint32_t>((i % 5) + 1);
doc_lens[i] = static_cast<uint32_t>(200 + i);
}
std::string encoded = BitPackedPostingList::encode(
doc_ids.data(), tfs.data(), doc_lens.data(), count, count, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
EXPECT_EQ(iter.cost(), count);
for (size_t i = 0; i < count; ++i) {
uint32_t doc = iter.next_doc();
EXPECT_EQ(doc, doc_ids[i]) << "Mismatch at index " << i;
EXPECT_EQ(iter.term_freq(), tfs[i]) << "TF mismatch at index " << i;
EXPECT_EQ(iter.doc_len(), doc_lens[i]) << "DocLen mismatch at index " << i;
}
EXPECT_EQ(iter.next_doc(), BitPackedPostingIterator::NO_MORE_DOCS);
}
// ============================================================
// Encode / Decode: large list (multiple blocks)
// ============================================================
TEST(BitPackedPostingListTest, EncodeDecodeLargeList) {
BM25Scorer scorer = make_scorer(10000, 500000);
const size_t count = 1000;
std::vector<uint32_t> doc_ids(count);
std::vector<uint32_t> tfs(count);
std::vector<uint32_t> doc_lens(count);
for (size_t i = 0; i < count; ++i) {
doc_ids[i] = static_cast<uint32_t>(i * 10);
tfs[i] = static_cast<uint32_t>((i % 20) + 1);
doc_lens[i] = static_cast<uint32_t>(50 + (i % 200));
}
std::string encoded = BitPackedPostingList::encode(
doc_ids.data(), tfs.data(), doc_lens.data(), count, count, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
EXPECT_EQ(iter.cost(), count);
for (size_t i = 0; i < count; ++i) {
uint32_t doc = iter.next_doc();
EXPECT_EQ(doc, doc_ids[i]) << "Mismatch at index " << i;
EXPECT_EQ(iter.term_freq(), tfs[i]) << "TF mismatch at index " << i;
EXPECT_EQ(iter.doc_len(), doc_lens[i]) << "DocLen mismatch at index " << i;
}
EXPECT_EQ(iter.next_doc(), BitPackedPostingIterator::NO_MORE_DOCS);
}
// ============================================================
// advance(): basic skip-list functionality
// ============================================================
TEST(BitPackedPostingListTest, AdvanceToExactDocId) {
BM25Scorer scorer = make_scorer();
const size_t count = 500;
std::vector<uint32_t> doc_ids(count);
std::vector<uint32_t> tfs(count, 1);
std::vector<uint32_t> doc_lens(count, 100);
for (size_t i = 0; i < count; ++i) {
doc_ids[i] = static_cast<uint32_t>(i * 3);
}
std::string encoded = BitPackedPostingList::encode(
doc_ids.data(), tfs.data(), doc_lens.data(), count, count, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
// Advance to exact doc_id
EXPECT_EQ(iter.advance(300), 300u);
EXPECT_EQ(iter.doc_id(), 300u);
// Advance to a doc_id that doesn't exist (should return next >= target)
EXPECT_EQ(iter.advance(301), 303u);
EXPECT_EQ(iter.doc_id(), 303u);
}
TEST(BitPackedPostingListTest, AdvanceToFirstDoc) {
BM25Scorer scorer = make_scorer();
uint32_t doc_ids[] = {10, 20, 30, 40, 50};
uint32_t tfs[] = {1, 2, 3, 4, 5};
uint32_t doc_lens[] = {100, 200, 300, 400, 500};
std::string encoded =
BitPackedPostingList::encode(doc_ids, tfs, doc_lens, 5, 5, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
// Advance to 0 should return the first doc (10)
EXPECT_EQ(iter.advance(0), 10u);
EXPECT_EQ(iter.term_freq(), 1u);
EXPECT_EQ(iter.doc_len(), 100u);
}
TEST(BitPackedPostingListTest, AdvanceBeyondLastDoc) {
BM25Scorer scorer = make_scorer();
uint32_t doc_ids[] = {10, 20, 30};
uint32_t tfs[] = {1, 2, 3};
uint32_t doc_lens[] = {100, 200, 300};
std::string encoded =
BitPackedPostingList::encode(doc_ids, tfs, doc_lens, 3, 3, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
EXPECT_EQ(iter.advance(31), BitPackedPostingIterator::NO_MORE_DOCS);
}
TEST(BitPackedPostingListTest, AdvanceAcrossBlocks) {
BM25Scorer scorer = make_scorer();
const size_t count = 300;
std::vector<uint32_t> doc_ids(count);
std::vector<uint32_t> tfs(count, 2);
std::vector<uint32_t> doc_lens(count, 50);
for (size_t i = 0; i < count; ++i) {
doc_ids[i] = static_cast<uint32_t>(i * 5);
}
std::string encoded = BitPackedPostingList::encode(
doc_ids.data(), tfs.data(), doc_lens.data(), count, count, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
// Advance from start to a doc in the 3rd block (block 2, index 256+)
// Block 0: doc_ids 0..635 (indices 0..127)
// Block 1: doc_ids 640..1275 (indices 128..255)
// Block 2: doc_ids 1280..1495 (indices 256..299)
EXPECT_EQ(iter.advance(1280), 1280u);
EXPECT_EQ(iter.doc_id(), 1280u);
EXPECT_EQ(iter.term_freq(), 2u);
// Continue with next_doc
EXPECT_EQ(iter.next_doc(), 1285u);
}
TEST(BitPackedPostingListTest, AdvanceSequentialCalls) {
BM25Scorer scorer = make_scorer();
const size_t count = 200;
std::vector<uint32_t> doc_ids(count);
std::vector<uint32_t> tfs(count, 1);
std::vector<uint32_t> doc_lens(count, 100);
for (size_t i = 0; i < count; ++i) {
doc_ids[i] = static_cast<uint32_t>(i * 7);
}
std::string encoded = BitPackedPostingList::encode(
doc_ids.data(), tfs.data(), doc_lens.data(), count, count, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
// Multiple sequential advance calls
EXPECT_EQ(iter.advance(100), 105u); // 15*7=105
EXPECT_EQ(iter.advance(500), 504u); // 72*7=504
EXPECT_EQ(iter.advance(1000), 1001u); // 143*7=1001
EXPECT_EQ(iter.advance(1400), BitPackedPostingIterator::NO_MORE_DOCS);
}
// ============================================================
// advance() after next_doc()
// ============================================================
TEST(BitPackedPostingListTest, AdvanceAfterNextDoc) {
BM25Scorer scorer = make_scorer();
const size_t count = 256;
std::vector<uint32_t> doc_ids(count);
std::vector<uint32_t> tfs(count, 1);
std::vector<uint32_t> doc_lens(count, 50);
for (size_t i = 0; i < count; ++i) {
doc_ids[i] = static_cast<uint32_t>(i * 4);
}
std::string encoded = BitPackedPostingList::encode(
doc_ids.data(), tfs.data(), doc_lens.data(), count, count, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
// Read a few docs
EXPECT_EQ(iter.next_doc(), 0u);
EXPECT_EQ(iter.next_doc(), 4u);
EXPECT_EQ(iter.next_doc(), 8u);
// Now advance past the current block
EXPECT_EQ(iter.advance(600), 600u); // 150*4=600
EXPECT_EQ(iter.term_freq(), 1u);
// Continue with next_doc
EXPECT_EQ(iter.next_doc(), 604u);
}
// ============================================================
// block_max_score correctness
// ============================================================
TEST(BitPackedPostingListTest, BlockMaxScoreCorrectness) {
BM25Scorer scorer = make_scorer(100, 5000);
const size_t count = 256; // 2 blocks
std::vector<uint32_t> doc_ids(count);
std::vector<uint32_t> tfs(count);
std::vector<uint32_t> doc_lens(count);
for (size_t i = 0; i < count; ++i) {
doc_ids[i] = static_cast<uint32_t>(i);
tfs[i] = static_cast<uint32_t>((i % 10) + 1);
doc_lens[i] = static_cast<uint32_t>(50 + (i % 50));
}
std::string encoded = BitPackedPostingList::encode(
doc_ids.data(), tfs.data(), doc_lens.data(), count, count, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
// Verify block_max_score for block 0 via block_max_info_for()
auto info0 = iter.block_max_info_for(0);
// Manually compute max score for block 0
float expected_max = 0.0f;
for (size_t i = 0; i < 128; ++i) {
float s = scorer.score(count, tfs[i], doc_lens[i]);
expected_max = std::max(expected_max, s);
}
EXPECT_FLOAT_EQ(info0.block_max_score, expected_max);
EXPECT_EQ(info0.block_last_doc, 127u);
// Verify block_max_score for block 1 via block_max_info_for()
auto info1 = iter.block_max_info_for(128);
expected_max = 0.0f;
for (size_t i = 128; i < 256; ++i) {
float s = scorer.score(count, tfs[i], doc_lens[i]);
expected_max = std::max(expected_max, s);
}
EXPECT_FLOAT_EQ(info1.block_max_score, expected_max);
EXPECT_EQ(info1.block_last_doc, 255u);
}
// ============================================================
// max_score() (global)
// ============================================================
TEST(BitPackedPostingListTest, GlobalMaxScore) {
BM25Scorer scorer = make_scorer(100, 5000);
const size_t count = 256;
std::vector<uint32_t> doc_ids(count);
std::vector<uint32_t> tfs(count);
std::vector<uint32_t> doc_lens(count);
for (size_t i = 0; i < count; ++i) {
doc_ids[i] = static_cast<uint32_t>(i);
tfs[i] = static_cast<uint32_t>((i % 10) + 1);
doc_lens[i] = static_cast<uint32_t>(50 + (i % 50));
}
std::string encoded = BitPackedPostingList::encode(
doc_ids.data(), tfs.data(), doc_lens.data(), count, count, scorer);
BitPackedPostingIterator iter;
EXPECT_EQ(iter.open(encoded.data(), encoded.size()), 0);
// Global max_score should be the maximum of all block_max_scores
float global_max = 0.0f;
for (size_t i = 0; i < count; ++i) {
float s = scorer.score(count, tfs[i], doc_lens[i]);
global_max = std::max(global_max, s);
}
EXPECT_FLOAT_EQ(iter.max_score(), global_max);
}
// ============================================================
// is_bitpacked_format()
// ============================================================
TEST(BitPackedPostingListTest, IsBitpackedFormatTrue) {
BM25Scorer scorer = make_scorer();
uint32_t doc_ids[] = {1};
uint32_t tfs[] = {1};
uint32_t doc_lens[] = {10};
std::string encoded =
BitPackedPostingList::encode(doc_ids, tfs, doc_lens, 1, 1, scorer);
EXPECT_TRUE(BitPackedPostingList::is_bitpacked_format(encoded.data(),
encoded.size()));
}
TEST(BitPackedPostingListTest, IsBitpackedFormatFalse) {
// Random data that doesn't start with the magic number
std::string random_data = "hello world";
EXPECT_FALSE(BitPackedPostingList::is_bitpacked_format(random_data.data(),
random_data.size()));
}
TEST(BitPackedPostingListTest, IsBitpackedFormatTooShort) {
std::string short_data = "ab";
EXPECT_FALSE(BitPackedPostingList::is_bitpacked_format(short_data.data(),
short_data.size()));
}
// ============================================================
// Error handling: open() with invalid data
// ============================================================
TEST(BitPackedPostingListTest, OpenWithNullData) {
BitPackedPostingIterator iter;
EXPECT_NE(iter.open(nullptr, 0), 0);
}
TEST(BitPackedPostingListTest, OpenWithTruncatedHeader) {
BitPackedPostingIterator iter;
char data[4] = {0};
EXPECT_NE(iter.open(data, 4), 0);
}
TEST(BitPackedPostingListTest, OpenWithBadMagic) {
BitPackedPostingIterator iter;
char data[16] = {0};
EXPECT_NE(iter.open(data, 16), 0);
}
// ============================================================
// Cross-path encode/decode: verify that posting lists encoded via the
// dispatch path (SIMD on x86) are correctly decodable by the scalar
// fallback. This guards against the full encode() pipeline drifting
// from the scalar decoder — the low-level ScalarLayoutMatchesSimd test
// only covers the pack/unpack primitives, not the complete block
// layout produced by encode().
// ============================================================
TEST(BitPackedPostingListTest, EncodeDecodeScalarCrossDecode) {
BM25Scorer scorer = make_scorer(1000, 50000);
// 300 docs → 2 full blocks (128 each) + 1 tail block (44)
const size_t count = 300;
std::vector<uint32_t> doc_ids(count);
std::vector<uint32_t> tfs(count);
std::vector<uint32_t> doc_lens(count);
for (size_t i = 0; i < count; ++i) {
doc_ids[i] = static_cast<uint32_t>(i * 7 + 3);
tfs[i] = static_cast<uint32_t>((i % 15) + 1);
doc_lens[i] = static_cast<uint32_t>(30 + (i % 100));
}
std::string encoded = BitPackedPostingList::encode(
doc_ids.data(), tfs.data(), doc_lens.data(), count, count, scorer);
// Parse the header
BitPackedPostingList::Header hdr{};
std::memcpy(&hdr, encoded.data(), sizeof(hdr));
ASSERT_EQ(hdr.magic, BitPackedPostingList::MAGIC);
ASSERT_EQ(hdr.num_docs, count);
const auto *skip = reinterpret_cast<const BitPackedPostingList::BlockMeta *>(
encoded.data() + BitPackedPostingList::HEADER_SIZE);
// Manually decode each block using the scalar path and verify
std::vector<uint32_t> decoded_doc_ids;
for (uint32_t b = 0; b < hdr.num_blocks; ++b) {
const char *block_ptr = encoded.data() + skip[b].block_offset;
BitPackedPostingList::BlockHeader bhdr{};
std::memcpy(&bhdr, block_ptr, sizeof(bhdr));
const uint8_t *packed =
reinterpret_cast<const uint8_t *>(block_ptr + sizeof(bhdr));
const bool is_full =
(bhdr.num_docs == BitPackedPostingList::DOCS_PER_BLOCK);
// Decode doc_id deltas using explicit scalar path
alignas(16) uint32_t deltas[BitPackedPostingList::DOCS_PER_BLOCK]{};
if (is_full) {
simd::scalar_unpack_uint32_128(packed, bhdr.bitwidth_id, deltas);
} else {
BitPackedPostingList::unpack_uint32(packed, bhdr.bitwidth_id,
bhdr.num_docs, deltas);
}
// Reconstruct absolute doc_ids via prefix-sum
uint32_t prev = bhdr.min_doc_id;
decoded_doc_ids.push_back(prev);
for (uint32_t i = 1; i < bhdr.num_docs; ++i) {
prev += deltas[i];
decoded_doc_ids.push_back(prev);
}
const size_t id_bytes =
is_full ? BitPackedPostingList::simd_packed_byte_size(bhdr.bitwidth_id)
: BitPackedPostingList::packed_byte_size(bhdr.bitwidth_id,
bhdr.num_docs);
const uint8_t *tf_packed = packed + id_bytes;
alignas(16) uint32_t decoded_tfs[BitPackedPostingList::DOCS_PER_BLOCK]{};
if (is_full) {
simd::scalar_unpack_uint32_128(tf_packed, bhdr.bitwidth_tf, decoded_tfs);
} else {
BitPackedPostingList::unpack_uint32(tf_packed, bhdr.bitwidth_tf,
bhdr.num_docs, decoded_tfs);
}
const size_t tf_bytes =
is_full ? BitPackedPostingList::simd_packed_byte_size(bhdr.bitwidth_tf)
: BitPackedPostingList::packed_byte_size(bhdr.bitwidth_tf,
bhdr.num_docs);
const uint8_t *dl_packed = tf_packed + tf_bytes;
alignas(16) uint32_t decoded_dls[BitPackedPostingList::DOCS_PER_BLOCK]{};
if (is_full) {
simd::scalar_unpack_uint32_128(dl_packed, bhdr.bitwidth_dl, decoded_dls);
} else {
BitPackedPostingList::unpack_uint32(dl_packed, bhdr.bitwidth_dl,
bhdr.num_docs, decoded_dls);
}
const size_t start =
static_cast<size_t>(b) * BitPackedPostingList::DOCS_PER_BLOCK;
for (uint32_t i = 0; i < bhdr.num_docs; ++i) {
EXPECT_EQ(decoded_tfs[i], tfs[start + i])
<< "TF mismatch block " << b << " index " << i;
EXPECT_EQ(decoded_dls[i], doc_lens[start + i])
<< "DocLen mismatch block " << b << " index " << i;
}
}
ASSERT_EQ(decoded_doc_ids.size(), count);
for (size_t i = 0; i < count; ++i) {
EXPECT_EQ(decoded_doc_ids[i], doc_ids[i])
<< "DocId mismatch at index " << i;
}
}
// ============================================================
// Consistency: advance() vs sequential next_doc()
// ============================================================
TEST(BitPackedPostingListTest, AdvanceConsistentWithNextDoc) {
BM25Scorer scorer = make_scorer();
const size_t count = 500;
std::vector<uint32_t> doc_ids(count);
std::vector<uint32_t> tfs(count);
std::vector<uint32_t> doc_lens(count);
std::mt19937 rng(42);
uint32_t current = 0;
for (size_t i = 0; i < count; ++i) {
current += (rng() % 10) + 1;
doc_ids[i] = current;
tfs[i] = (rng() % 10) + 1;
doc_lens[i] = (rng() % 200) + 10;
}
std::string encoded = BitPackedPostingList::encode(
doc_ids.data(), tfs.data(), doc_lens.data(), count, count, scorer);
// Collect all docs via next_doc
BitPackedPostingIterator iter1;
EXPECT_EQ(iter1.open(encoded.data(), encoded.size()), 0);
std::vector<uint32_t> all_docs;
std::vector<uint32_t> all_tfs;
std::vector<uint32_t> all_doc_lens;
uint32_t doc = iter1.next_doc();
while (doc != BitPackedPostingIterator::NO_MORE_DOCS) {
all_docs.push_back(doc);
all_tfs.push_back(iter1.term_freq());
all_doc_lens.push_back(iter1.doc_len());
doc = iter1.next_doc();
}
ASSERT_EQ(all_docs.size(), count);
// Verify advance to various targets matches sequential scan
BitPackedPostingIterator iter2;
EXPECT_EQ(iter2.open(encoded.data(), encoded.size()), 0);
std::vector<uint32_t> targets = {0,
1,
doc_ids[50],
doc_ids[127],
doc_ids[128],
doc_ids[200],
doc_ids[count - 1]};
for (uint32_t target : targets) {
BitPackedPostingIterator iter_adv;
EXPECT_EQ(iter_adv.open(encoded.data(), encoded.size()), 0);
uint32_t adv_doc = iter_adv.advance(target);
// Find expected result via linear scan
auto it = std::lower_bound(all_docs.begin(), all_docs.end(), target);
if (it == all_docs.end()) {
EXPECT_EQ(adv_doc, BitPackedPostingIterator::NO_MORE_DOCS)
<< "target=" << target;
} else {
size_t idx = it - all_docs.begin();
EXPECT_EQ(adv_doc, all_docs[idx]) << "target=" << target;
EXPECT_EQ(iter_adv.term_freq(), all_tfs[idx]) << "target=" << target;
EXPECT_EQ(iter_adv.doc_len(), all_doc_lens[idx]) << "target=" << target;
}
}
}
@@ -0,0 +1,560 @@
// 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 "db/index/column/fts_column/fts_ast_rewriter.h"
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include <gtest/gtest.h>
#include "db/index/column/fts_column/fts_query_ast.h"
namespace zvec::fts {
namespace {
// Convenience constructors keep the test bodies focused on what's being
// asserted rather than on AST scaffolding.
FtsAstNodePtr term(const std::string &t, bool must = false,
bool must_not = false, float boost = 1.0f) {
auto n = std::make_unique<TermNode>(t, must, must_not);
n->boost = boost;
return n;
}
FtsAstNodePtr phrase(std::vector<std::string> ts, bool must = false,
bool must_not = false, float boost = 1.0f) {
auto n = std::make_unique<PhraseNode>();
n->terms = std::move(ts);
n->must = must;
n->must_not = must_not;
n->boost = boost;
return n;
}
FtsAstNodePtr empty_node() {
return std::make_unique<EmptyNode>();
}
template <typename Node>
FtsAstNodePtr composite(std::vector<FtsAstNodePtr> children, bool must = false,
bool must_not = false) {
auto n = std::make_unique<Node>();
n->children = std::move(children);
n->must = must;
n->must_not = must_not;
return n;
}
FtsAstNodePtr or_node(std::vector<FtsAstNodePtr> c, bool must = false,
bool must_not = false) {
return composite<OrNode>(std::move(c), must, must_not);
}
FtsAstNodePtr and_node(std::vector<FtsAstNodePtr> c, bool must = false,
bool must_not = false) {
return composite<AndNode>(std::move(c), must, must_not);
}
// Pull the single TermNode child out of a composite for boost assertions.
const TermNode &as_term(const FtsAstNode &n) {
return static_cast<const TermNode &>(n);
}
const PhraseNode &as_phrase(const FtsAstNode &n) {
return static_cast<const PhraseNode &>(n);
}
const OrNode &as_or(const FtsAstNode &n) {
return static_cast<const OrNode &>(n);
}
const AndNode &as_and(const FtsAstNode &n) {
return static_cast<const AndNode &>(n);
}
} // namespace
// --- Dedup ---
TEST(FtsAstRewriterTest, OrDedupsRepeatedTerms) {
// OR(apple, apple, banana) → OR(apple^2, banana)
std::vector<FtsAstNodePtr> children;
children.push_back(term("apple"));
children.push_back(term("apple"));
children.push_back(term("banana"));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
const auto &n = as_or(*ast);
ASSERT_EQ(n.children.size(), 2u);
EXPECT_EQ(as_term(*n.children[0]).term, "apple");
EXPECT_FLOAT_EQ(n.children[0]->boost, 2.0f);
EXPECT_EQ(as_term(*n.children[1]).term, "banana");
EXPECT_FLOAT_EQ(n.children[1]->boost, 1.0f);
}
TEST(FtsAstRewriterTest, AndDedupsRepeatedTerms) {
std::vector<FtsAstNodePtr> children;
children.push_back(term("apple"));
children.push_back(term("apple"));
children.push_back(term("apple"));
auto ast = and_node(std::move(children));
simplify(ast);
// Single-child fold collapses AND(apple^3) → apple^3.
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_EQ(as_term(*ast).term, "apple");
EXPECT_FLOAT_EQ(ast->boost, 3.0f);
}
TEST(FtsAstRewriterTest, DifferentOccurDoesNotMerge) {
// OR(apple, +apple) — the +apple becomes a must clause, plain apple becomes
// should. After canonicalization into AND(apple, ?apple) and dedup (same term
// text, same must/must_not bits post-canonicalization), they merge into a
// single term with boost=2.0.
std::vector<FtsAstNodePtr> children;
children.push_back(term("apple"));
children.push_back(term("apple", /*must=*/true));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_FLOAT_EQ(ast->boost, 2.0f);
}
// --- Conflict ---
TEST(FtsAstRewriterTest, AndMustVsMustNotSameTermBecomesEmpty) {
std::vector<FtsAstNodePtr> children;
children.push_back(term("apple", /*must=*/true));
children.push_back(term("apple", /*must=*/false, /*must_not=*/true));
children.push_back(term("banana"));
auto ast = and_node(std::move(children));
simplify(ast);
EXPECT_EQ(ast->type(), FtsNodeType::EMPTY);
}
TEST(FtsAstRewriterTest, AndAllMustNotBecomesEmpty) {
std::vector<FtsAstNodePtr> children;
children.push_back(term("apple", false, true));
children.push_back(term("banana", false, true));
auto ast = and_node(std::move(children));
simplify(ast);
EXPECT_EQ(ast->type(), FtsNodeType::EMPTY);
}
// --- Flattening ---
TEST(FtsAstRewriterTest, OrFlattensNestedOr) {
// OR(a, OR(b, c)) → OR(a, b, c)
std::vector<FtsAstNodePtr> inner;
inner.push_back(term("b"));
inner.push_back(term("c"));
std::vector<FtsAstNodePtr> outer;
outer.push_back(term("a"));
outer.push_back(or_node(std::move(inner)));
auto ast = or_node(std::move(outer));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
ASSERT_EQ(as_or(*ast).children.size(), 3u);
}
TEST(FtsAstRewriterTest, AndFlattensNestedAndWithoutMustNot) {
std::vector<FtsAstNodePtr> inner;
inner.push_back(term("b"));
inner.push_back(term("c", true));
std::vector<FtsAstNodePtr> outer;
outer.push_back(term("a"));
outer.push_back(and_node(std::move(inner)));
auto ast = and_node(std::move(outer));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
EXPECT_EQ(as_and(*ast).children.size(), 3u);
}
TEST(FtsAstRewriterTest, AndDoesNotFlattenAndWithMustNotChild) {
std::vector<FtsAstNodePtr> inner;
inner.push_back(term("b"));
inner.push_back(term("c", false, true));
std::vector<FtsAstNodePtr> outer;
outer.push_back(term("a"));
outer.push_back(and_node(std::move(inner)));
auto ast = and_node(std::move(outer));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
EXPECT_EQ(as_and(*ast).children.size(), 2u);
}
TEST(FtsAstRewriterTest, FlattenThenDedupCrossLayer) {
// OR(a, OR(a, b)) → flatten → OR(a, a, b) → dedup → OR(a^2, b)
std::vector<FtsAstNodePtr> inner;
inner.push_back(term("a"));
inner.push_back(term("b"));
std::vector<FtsAstNodePtr> outer;
outer.push_back(term("a"));
outer.push_back(or_node(std::move(inner)));
auto ast = or_node(std::move(outer));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
const auto &n = as_or(*ast);
ASSERT_EQ(n.children.size(), 2u);
EXPECT_EQ(as_term(*n.children[0]).term, "a");
EXPECT_FLOAT_EQ(n.children[0]->boost, 2.0f);
EXPECT_EQ(as_term(*n.children[1]).term, "b");
}
// --- Phrase ---
TEST(FtsAstRewriterTest, PhraseSameTermsAreDeduped) {
std::vector<FtsAstNodePtr> children;
children.push_back(phrase({"new", "york"}));
children.push_back(phrase({"new", "york"}));
children.push_back(phrase({"new", "york"}));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::PHRASE);
EXPECT_FLOAT_EQ(ast->boost, 3.0f);
}
TEST(FtsAstRewriterTest, PhraseInternalRepeatNotMerged) {
// Position-sensitive: "new new york" must keep its internal duplication.
auto p = phrase({"new", "new", "york"});
FtsAstNodePtr ast = std::move(p);
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::PHRASE);
ASSERT_EQ(as_phrase(*ast).terms.size(), 3u);
EXPECT_EQ(as_phrase(*ast).terms[0], "new");
EXPECT_EQ(as_phrase(*ast).terms[1], "new");
EXPECT_EQ(as_phrase(*ast).terms[2], "york");
}
TEST(FtsAstRewriterTest, DifferentPhrasesDoNotMerge) {
std::vector<FtsAstNodePtr> children;
children.push_back(phrase({"new", "york"}));
children.push_back(phrase({"york", "new"}));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
EXPECT_EQ(as_or(*ast).children.size(), 2u);
}
// --- EmptyNode propagation ---
TEST(FtsAstRewriterTest, AndWithEmptyChildShortCircuits) {
std::vector<FtsAstNodePtr> children;
children.push_back(term("apple"));
children.push_back(empty_node());
auto ast = and_node(std::move(children));
simplify(ast);
EXPECT_EQ(ast->type(), FtsNodeType::EMPTY);
}
TEST(FtsAstRewriterTest, OrDropsEmptyChild) {
std::vector<FtsAstNodePtr> children;
children.push_back(term("apple"));
children.push_back(empty_node());
children.push_back(term("banana"));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
EXPECT_EQ(as_or(*ast).children.size(), 2u);
}
TEST(FtsAstRewriterTest, MustNotEmptyInAndIsNoOp) {
// AND(apple, -EMPTY) — excluding nothing has no effect.
std::vector<FtsAstNodePtr> children;
children.push_back(term("apple"));
auto e = std::make_unique<EmptyNode>();
e->must_not = true;
children.push_back(std::move(e));
auto ast = and_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_EQ(as_term(*ast).term, "apple");
}
// --- Single-child fold ---
TEST(FtsAstRewriterTest, SingleChildOrFolds) {
std::vector<FtsAstNodePtr> children;
children.push_back(term("apple"));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_EQ(as_term(*ast).term, "apple");
}
TEST(FtsAstRewriterTest, FoldedSingleChildInheritsParentModifier) {
// +OR(apple) → +apple (must flag lifts onto the surviving child)
std::vector<FtsAstNodePtr> children;
children.push_back(term("apple"));
auto ast = or_node(std::move(children), /*must=*/true);
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_TRUE(ast->must);
EXPECT_FALSE(ast->must_not);
}
// --- Idempotence ---
TEST(FtsAstRewriterTest, SimplifyIsIdempotent) {
// Build something gnarly enough to exercise multiple rules at once.
std::vector<FtsAstNodePtr> inner_or;
inner_or.push_back(term("a"));
inner_or.push_back(term("a"));
std::vector<FtsAstNodePtr> children;
children.push_back(term("a"));
children.push_back(or_node(std::move(inner_or)));
children.push_back(term("b"));
children.push_back(empty_node());
auto ast = or_node(std::move(children));
simplify(ast);
const std::string after_first = ast->text();
simplify(ast);
const std::string after_second = ast->text();
EXPECT_EQ(after_first, after_second);
}
// --- OR must_not canonicalization ---
TEST(FtsAstRewriterTest, OrWithSinglePositiveAndMustNotBecomesAnd) {
// OR(a, -b) → AND(a, -b)
std::vector<FtsAstNodePtr> children;
children.push_back(term("a"));
children.push_back(term("b", false, true));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
const auto &n = as_and(*ast);
ASSERT_EQ(n.children.size(), 2u);
EXPECT_EQ(as_term(*n.children[0]).term, "a");
EXPECT_FALSE(n.children[0]->must_not);
EXPECT_EQ(as_term(*n.children[1]).term, "b");
EXPECT_TRUE(n.children[1]->must_not);
}
TEST(FtsAstRewriterTest, OrWithMultiplePositivesAndMustNotWrapsInAnd) {
// OR(a, b, -c) → AND(OR(a, b), -c)
std::vector<FtsAstNodePtr> children;
children.push_back(term("a"));
children.push_back(term("b"));
children.push_back(term("c", false, true));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
const auto &n = as_and(*ast);
ASSERT_EQ(n.children.size(), 2u);
ASSERT_EQ(n.children[0]->type(), FtsNodeType::OR);
EXPECT_EQ(as_or(*n.children[0]).children.size(), 2u);
EXPECT_EQ(as_term(*n.children[1]).term, "c");
EXPECT_TRUE(n.children[1]->must_not);
}
TEST(FtsAstRewriterTest, OrCanonicalizationCatchesSameTermConflict) {
// OR(a, -a) — canonicalization moves -a into AND with a, then
// and_has_mustnot_conflict fires → EmptyNode.
std::vector<FtsAstNodePtr> children;
children.push_back(term("a"));
children.push_back(term("a", false, true));
auto ast = or_node(std::move(children));
simplify(ast);
EXPECT_EQ(ast->type(), FtsNodeType::EMPTY);
}
TEST(FtsAstRewriterTest, OrCanonicalizationLiftsParentModifier) {
// +OR(a, -b) → +AND(a, -b)
std::vector<FtsAstNodePtr> children;
children.push_back(term("a"));
children.push_back(term("b", false, true));
auto ast = or_node(std::move(children), /*must=*/true);
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
EXPECT_TRUE(ast->must);
EXPECT_FALSE(ast->must_not);
}
TEST(FtsAstRewriterTest, NestedOrWithMustNotCanonicalizedAtBothLevels) {
// OR(x, OR(b, -c)) — inner canonicalizes to AND(b, -c); outer keeps OR
// since it has no must_not after recursion.
std::vector<FtsAstNodePtr> inner;
inner.push_back(term("b"));
inner.push_back(term("c", false, true));
std::vector<FtsAstNodePtr> outer;
outer.push_back(term("x"));
outer.push_back(or_node(std::move(inner)));
auto ast = or_node(std::move(outer));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
const auto &n = as_or(*ast);
ASSERT_EQ(n.children.size(), 2u);
EXPECT_EQ(as_term(*n.children[0]).term, "x");
EXPECT_EQ(n.children[1]->type(), FtsNodeType::AND);
}
TEST(FtsAstRewriterTest, OrWithoutMustNotIsLeftAlone) {
std::vector<FtsAstNodePtr> children;
children.push_back(term("a"));
children.push_back(term("b"));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
EXPECT_EQ(as_or(*ast).children.size(), 2u);
}
// --- OR must canonicalization ---
TEST(FtsAstRewriterTest, OrWithMustChildrenCanonicalizesToAnd) {
// OR(foo, +bar, baz, +bay) → AND(bar, bay, ?OR(foo, baz))
std::vector<FtsAstNodePtr> children;
children.push_back(term("foo"));
children.push_back(term("bar", /*must=*/true));
children.push_back(term("baz"));
children.push_back(term("bay", /*must=*/true));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
const auto &n = as_and(*ast);
ASSERT_EQ(n.children.size(), 3u);
// First two children: bar, bay (must terms, must flag cleared)
EXPECT_EQ(as_term(*n.children[0]).term, "bar");
EXPECT_FALSE(n.children[0]->must);
EXPECT_FALSE(n.children[0]->should);
EXPECT_EQ(as_term(*n.children[1]).term, "bay");
EXPECT_FALSE(n.children[1]->must);
EXPECT_FALSE(n.children[1]->should);
// Third child: ?OR(foo, baz) with should=true
ASSERT_EQ(n.children[2]->type(), FtsNodeType::OR);
EXPECT_TRUE(n.children[2]->should);
const auto &should_or = as_or(*n.children[2]);
ASSERT_EQ(should_or.children.size(), 2u);
EXPECT_EQ(as_term(*should_or.children[0]).term, "foo");
EXPECT_EQ(as_term(*should_or.children[1]).term, "baz");
}
TEST(FtsAstRewriterTest, OrWithSingleMustAndSingleShouldFoldsCorrectly) {
// OR(foo, +bar) → AND(bar, ?foo)
std::vector<FtsAstNodePtr> children;
children.push_back(term("foo"));
children.push_back(term("bar", /*must=*/true));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
const auto &n = as_and(*ast);
ASSERT_EQ(n.children.size(), 2u);
EXPECT_EQ(as_term(*n.children[0]).term, "bar");
EXPECT_FALSE(n.children[0]->should);
EXPECT_EQ(as_term(*n.children[1]).term, "foo");
EXPECT_TRUE(n.children[1]->should);
}
TEST(FtsAstRewriterTest, OrWithAllMustNoShouldChildren) {
// OR(+bar, +bay) → AND(bar, bay) — no should part
std::vector<FtsAstNodePtr> children;
children.push_back(term("bar", /*must=*/true));
children.push_back(term("bay", /*must=*/true));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
const auto &n = as_and(*ast);
ASSERT_EQ(n.children.size(), 2u);
EXPECT_EQ(as_term(*n.children[0]).term, "bar");
EXPECT_EQ(as_term(*n.children[1]).term, "bay");
}
TEST(FtsAstRewriterTest, OrWithMustAndMustNotCanonicalizesCorrectly) {
// OR(foo, +bar, -baz) → must_not is processed first, then must
std::vector<FtsAstNodePtr> children;
children.push_back(term("foo"));
children.push_back(term("bar", /*must=*/true));
children.push_back(term("baz", /*must=*/false, /*must_not=*/true));
auto ast = or_node(std::move(children));
simplify(ast);
// Should become AND-like structure with bar required, baz excluded
ASSERT_EQ(ast->type(), FtsNodeType::AND);
}
TEST(FtsAstRewriterTest, OrWithSingleMustNoShouldFoldsToTerm) {
// OR(+bar) → bar (single-child fold after must canonicalization)
std::vector<FtsAstNodePtr> children;
children.push_back(term("bar", /*must=*/true));
auto ast = or_node(std::move(children));
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_EQ(as_term(*ast).term, "bar");
}
// --- Leaf untouched ---
TEST(FtsAstRewriterTest, BareTermPassthrough) {
FtsAstNodePtr ast = term("apple");
simplify(ast);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_EQ(as_term(*ast).term, "apple");
EXPECT_FLOAT_EQ(ast->boost, 1.0f);
}
} // namespace zvec::fts
@@ -0,0 +1,98 @@
// 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 "db/index/column/fts_column/iterator/fts_candidate_iterator.h"
#include <cstdint>
#include <vector>
#include <gtest/gtest.h>
#include "db/index/column/fts_column/iterator/fts_doc_iterator.h"
using zvec::fts::CandidateDocIterator;
using zvec::fts::DocIterator;
namespace {
constexpr uint32_t kNoMore = DocIterator::NO_MORE_DOCS;
} // namespace
TEST(CandidateDocIteratorTest, EmptyVectorYieldsNothing) {
CandidateDocIterator it({});
EXPECT_EQ(it.cost(), 0u);
EXPECT_EQ(it.next_doc(), kNoMore);
EXPECT_EQ(it.next_doc(), kNoMore);
}
TEST(CandidateDocIteratorTest, NextDocStreamsAscending) {
CandidateDocIterator it({0, 5, 10, 100});
EXPECT_EQ(it.cost(), 4u);
EXPECT_FLOAT_EQ(it.max_score(), 0.0f);
EXPECT_FLOAT_EQ(it.score(), 0.0f);
EXPECT_TRUE(it.matches());
EXPECT_EQ(it.next_doc(), 0u);
EXPECT_EQ(it.doc_id(), 0u);
EXPECT_EQ(it.next_doc(), 5u);
EXPECT_EQ(it.next_doc(), 10u);
EXPECT_EQ(it.next_doc(), 100u);
EXPECT_EQ(it.next_doc(), kNoMore);
EXPECT_EQ(it.next_doc(), kNoMore);
}
TEST(CandidateDocIteratorTest, AdvanceLandsOnExactMatch) {
CandidateDocIterator it({10, 20, 30, 40, 50});
EXPECT_EQ(it.advance(20), 20u);
EXPECT_EQ(it.doc_id(), 20u);
// Subsequent next_doc continues past the advanced position.
EXPECT_EQ(it.next_doc(), 30u);
}
TEST(CandidateDocIteratorTest, AdvanceSeeksToNextHigher) {
CandidateDocIterator it({10, 20, 30, 40, 50});
EXPECT_EQ(it.advance(25), 30u);
EXPECT_EQ(it.next_doc(), 40u);
}
TEST(CandidateDocIteratorTest, AdvancePastLastYieldsNoMore) {
CandidateDocIterator it({10, 20, 30});
EXPECT_EQ(it.advance(50), kNoMore);
EXPECT_EQ(it.next_doc(), kNoMore);
}
TEST(CandidateDocIteratorTest, AdvanceBeforeAnyConsumeWorks) {
CandidateDocIterator it({10, 20, 30});
EXPECT_EQ(it.advance(0), 10u);
EXPECT_EQ(it.next_doc(), 20u);
}
TEST(CandidateDocIteratorTest, AdvanceInterleavedWithNext) {
CandidateDocIterator it({5, 10, 15, 20, 25, 30});
EXPECT_EQ(it.next_doc(), 5u);
EXPECT_EQ(it.advance(15), 15u);
EXPECT_EQ(it.next_doc(), 20u);
EXPECT_EQ(it.advance(99), kNoMore);
}
TEST(CandidateDocIteratorTest, SingleElement) {
CandidateDocIterator it({42});
EXPECT_EQ(it.cost(), 1u);
EXPECT_EQ(it.advance(42), 42u);
EXPECT_EQ(it.next_doc(), kNoMore);
}
TEST(CandidateDocIteratorTest, AdvanceCachesDocId) {
CandidateDocIterator it({1, 2, 3});
EXPECT_EQ(it.advance(2), 2u);
EXPECT_EQ(it.doc_id(), 2u);
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,305 @@
// 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 <cstdint>
#include <map>
#include <memory>
#include <utility>
#include <vector>
#include <gtest/gtest.h>
#include "db/index/column/fts_column/iterator/fts_conjunction_iterator.h"
#include "db/index/column/fts_column/iterator/fts_disjunction_iterator.h"
#include "db/index/column/fts_column/iterator/fts_doc_iterator.h"
using zvec::fts::ConjunctionIterator;
using zvec::fts::DisjunctionIterator;
using zvec::fts::DocIterator;
using zvec::fts::DocIteratorPtr;
namespace {
constexpr uint32_t kNoMore = DocIterator::NO_MORE_DOCS;
// ---------------------------------------------------------------------------
// FakeDocIterator: controllable posting list with per-doc scores and
// per-block max score metadata.
// ---------------------------------------------------------------------------
class FakeDocIterator : public DocIterator {
public:
// Each entry: {doc_id, score}
// block_maxes: {block_last_doc -> block_max_score}
FakeDocIterator(std::vector<std::pair<uint32_t, float>> entries,
float max_score_val,
std::map<uint32_t, float> block_maxes = {})
: entries_(std::move(entries)),
max_score_val_(max_score_val),
block_maxes_(std::move(block_maxes)) {
cached_max_score_ = max_score_val_;
}
uint32_t next_doc() override {
++pos_;
if (pos_ >= entries_.size()) {
cached_doc_id_ = kNoMore;
return kNoMore;
}
cached_doc_id_ = entries_[pos_].first;
return cached_doc_id_;
}
uint32_t advance(uint32_t target) override {
while (pos_ + 1 < entries_.size() && entries_[pos_ + 1].first < target) {
++pos_;
}
return next_doc();
}
float score() override {
++score_call_count_;
if (pos_ < entries_.size()) {
return entries_[pos_].second;
}
return 0.0f;
}
uint64_t cost() const override {
return entries_.size();
}
float max_score() const override {
return max_score_val_;
}
BlockMaxInfo block_max_info_for(uint32_t target) const override {
if (block_maxes_.empty()) {
return {max_score_val_, kNoMore};
}
// Find the first block whose last_doc >= target
for (const auto &[last_doc, bm_score] : block_maxes_) {
if (last_doc >= target) {
return {bm_score, last_doc};
}
}
return {0.0f, kNoMore};
}
int score_call_count() const {
return score_call_count_;
}
private:
std::vector<std::pair<uint32_t, float>> entries_;
float max_score_val_;
std::map<uint32_t, float> block_maxes_;
size_t pos_{SIZE_MAX}; // before first element
int score_call_count_{0};
};
// Collect all doc_ids from an iterator
std::vector<uint32_t> collect_docs(DocIterator *iter) {
std::vector<uint32_t> docs;
uint32_t doc = iter->next_doc();
while (doc != kNoMore) {
if (iter->matches()) {
docs.push_back(doc);
}
doc = iter->next_doc();
}
return docs;
}
} // namespace
// ============================================================
// Optimization 1: Block-Max skip
// ============================================================
// Block 0 [0..127] block_max sum = 1.0 < threshold 2.0 → skipped
// Block 1 [128..255] block_max sum = 10.0 >= 2.0 → kept
TEST(ConjunctionOptTest, BlockMaxSkipNonCompetitiveBlocks) {
std::vector<std::pair<uint32_t, float>> list1 = {
{10, 0.3f}, {50, 0.4f}, {130, 4.0f}, {200, 3.5f}};
std::vector<std::pair<uint32_t, float>> list2 = {
{10, 0.2f}, {50, 0.3f}, {130, 5.0f}, {200, 4.0f}};
std::map<uint32_t, float> bm1 = {{127, 0.5f}, {255, 5.0f}};
std::map<uint32_t, float> bm2 = {{127, 0.5f}, {255, 5.0f}};
std::vector<DocIteratorPtr> musts;
musts.push_back(std::make_unique<FakeDocIterator>(list1, 5.0f, bm1));
musts.push_back(std::make_unique<FakeDocIterator>(list2, 5.0f, bm2));
ConjunctionIterator conj(std::move(musts), std::vector<DocIteratorPtr>{});
conj.set_min_competitive_score(2.0f);
auto docs = collect_docs(&conj);
ASSERT_EQ(docs.size(), 2u);
EXPECT_EQ(docs[0], 130u);
EXPECT_EQ(docs[1], 200u);
}
// 3 blocks, threshold skips block 0 and block 1, only block 2 survives
TEST(ConjunctionOptTest, BlockMaxSkipMultipleBlocks) {
std::vector<std::pair<uint32_t, float>> list1 = {
{10, 0.1f}, {130, 1.5f}, {260, 4.0f}};
std::vector<std::pair<uint32_t, float>> list2 = {
{10, 0.1f}, {130, 1.0f}, {260, 4.5f}};
std::map<uint32_t, float> bm1 = {{127, 0.5f}, {255, 2.0f}, {383, 5.0f}};
std::map<uint32_t, float> bm2 = {{127, 0.5f}, {255, 2.0f}, {383, 5.0f}};
std::vector<DocIteratorPtr> musts;
musts.push_back(std::make_unique<FakeDocIterator>(list1, 5.0f, bm1));
musts.push_back(std::make_unique<FakeDocIterator>(list2, 5.0f, bm2));
ConjunctionIterator conj(std::move(musts), std::vector<DocIteratorPtr>{});
// block 0 sum=1.0, block 1 sum=4.0, block 2 sum=10.0; threshold=5.0
conj.set_min_competitive_score(5.0f);
auto docs = collect_docs(&conj);
ASSERT_EQ(docs.size(), 1u);
EXPECT_EQ(docs[0], 260u);
}
// ============================================================
// Optimization 3: Score early-exit
// ============================================================
// 3 must iterators sorted by cost (ascending). After scoring the first
// (lowest-cost) iterator, accumulated + remaining_max < threshold →
// score() exits early, skipping the remaining two score() calls.
TEST(ConjunctionOptTest, ScoreEarlyExitReducesScoreCalls) {
// iter0: cost=1, max_score=2.0, actual_score=0.1
// iter1: cost=2, max_score=2.0, actual_score=1.0
// iter2: cost=3, max_score=2.0, actual_score=1.5
// cached_max_score_ = 6.0, threshold = 5.5
// After iter0: remaining_max=4.0, total=0.1. 0.1+4.0=4.1 < 5.5 → exit
auto *raw0 = new FakeDocIterator({{1, 0.1f}}, 2.0f);
auto *raw1 = new FakeDocIterator({{1, 1.0f}}, 2.0f);
auto *raw2 = new FakeDocIterator({{1, 1.5f}}, 2.0f);
std::vector<DocIteratorPtr> musts;
musts.emplace_back(raw0);
musts.emplace_back(raw1);
musts.emplace_back(raw2);
ConjunctionIterator conj(std::move(musts), std::vector<DocIteratorPtr>{});
conj.set_min_competitive_score(5.5f);
uint32_t doc = conj.next_doc();
ASSERT_EQ(doc, 1u);
conj.score();
// Only the first iterator's score() should have been called;
// the other two were skipped by early-exit.
EXPECT_EQ(raw0->score_call_count(), 1);
EXPECT_EQ(raw1->score_call_count(), 0);
EXPECT_EQ(raw2->score_call_count(), 0);
}
// When scores are competitive, all iterators' score() are called
TEST(ConjunctionOptTest, ScoreNoEarlyExitCallsAll) {
auto *raw0 = new FakeDocIterator({{1, 3.0f}}, 3.0f);
auto *raw1 = new FakeDocIterator({{1, 3.0f}}, 3.0f);
std::vector<DocIteratorPtr> musts;
musts.emplace_back(raw0);
musts.emplace_back(raw1);
ConjunctionIterator conj(std::move(musts), std::vector<DocIteratorPtr>{});
conj.set_min_competitive_score(5.0f);
uint32_t doc = conj.next_doc();
ASSERT_EQ(doc, 1u);
float s = conj.score();
EXPECT_FLOAT_EQ(s, 6.0f);
EXPECT_EQ(raw0->score_call_count(), 1);
EXPECT_EQ(raw1->score_call_count(), 1);
}
// ============================================================
// Optimization 4: Phrase threshold forwarding
// ============================================================
// set_min_competitive_score propagated into inner conjunction triggers
// block-max skip; without forwarding all docs would be returned.
TEST(ConjunctionOptTest, PhraseForwardingBlockMaxSkip) {
std::vector<std::pair<uint32_t, float>> list1 = {{10, 0.2f}, {130, 4.0f}};
std::vector<std::pair<uint32_t, float>> list2 = {{10, 0.2f}, {130, 5.0f}};
std::map<uint32_t, float> bm1 = {{127, 0.3f}, {255, 5.0f}};
std::map<uint32_t, float> bm2 = {{127, 0.3f}, {255, 5.0f}};
std::vector<DocIteratorPtr> musts;
musts.push_back(std::make_unique<FakeDocIterator>(list1, 5.0f, bm1));
musts.push_back(std::make_unique<FakeDocIterator>(list2, 5.0f, bm2));
auto inner = std::make_unique<ConjunctionIterator>(
std::move(musts), std::vector<DocIteratorPtr>{});
// Forward threshold as PhraseDocIterator would
inner->set_min_competitive_score(2.0f);
auto docs = collect_docs(inner.get());
// Block 0 (sum 0.6 < 2.0) skipped
ASSERT_EQ(docs.size(), 1u);
EXPECT_EQ(docs[0], 130u);
}
// ============================================================
// DisjunctionIterator::advance() bypass WAND fix
// ============================================================
// advance() must faithfully return target even with high min_competitive_score.
// Without the fix, advance() delegates to next_doc() which triggers WAND
// pruning and returns NO_MORE_DOCS.
TEST(ConjunctionOptTest, DisjunctionAdvanceBypassesWand) {
std::vector<DocIteratorPtr> sub_iters;
sub_iters.push_back(std::make_unique<FakeDocIterator>(
std::vector<std::pair<uint32_t, float>>{{1, 0.1f}, {5, 0.1f}, {10, 0.1f}},
1.0f));
sub_iters.push_back(std::make_unique<FakeDocIterator>(
std::vector<std::pair<uint32_t, float>>{{3, 0.2f}, {5, 0.2f}, {20, 0.2f}},
1.0f));
DisjunctionIterator disj(std::move(sub_iters));
disj.set_min_competitive_score(100.0f);
uint32_t doc = disj.advance(5);
EXPECT_EQ(doc, 5u);
EXPECT_TRUE(disj.matches());
doc = disj.advance(10);
EXPECT_EQ(doc, 10u);
}
TEST(ConjunctionOptTest, ConjunctionAdvanceBypassesCompetitivePruning) {
std::vector<std::pair<uint32_t, float>> list1 = {{5, 0.1f}, {130, 5.0f}};
std::vector<std::pair<uint32_t, float>> list2 = {{5, 0.1f}, {130, 5.0f}};
std::map<uint32_t, float> bm1 = {{127, 0.5f}, {255, 5.0f}};
std::map<uint32_t, float> bm2 = {{127, 0.5f}, {255, 5.0f}};
std::vector<DocIteratorPtr> musts;
musts.push_back(std::make_unique<FakeDocIterator>(list1, 5.0f, bm1));
musts.push_back(std::make_unique<FakeDocIterator>(list2, 5.0f, bm2));
ConjunctionIterator conj(std::move(musts), std::vector<DocIteratorPtr>{});
conj.set_min_competitive_score(2.0f);
uint32_t doc = conj.advance(5);
EXPECT_EQ(doc, 5u);
EXPECT_TRUE(conj.matches());
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,124 @@
// 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 <string>
#include <vector>
#include <gtest/gtest.h>
#include "db/index/column/fts_column/tokenizer/token_filter.h"
using namespace zvec::fts;
static std::vector<Token> make_tokens(const std::vector<std::string> &texts) {
std::vector<Token> tokens;
for (size_t i = 0; i < texts.size(); ++i) {
tokens.push_back({texts[i], 0, static_cast<uint32_t>(i)});
}
return tokens;
}
class LowercaseTokenFilterTest : public ::testing::Test {
protected:
LowercaseTokenFilter filter_;
};
TEST_F(LowercaseTokenFilterTest, AsciiBasic) {
auto tokens = make_tokens({"Hello", "WORLD", "FoO"});
auto result = filter_.filter(std::move(tokens));
ASSERT_EQ(result.size(), 3u);
EXPECT_EQ(result[0].text, "hello");
EXPECT_EQ(result[1].text, "world");
EXPECT_EQ(result[2].text, "foo");
}
TEST_F(LowercaseTokenFilterTest, AlreadyLowercase) {
auto tokens = make_tokens({"already", "lower"});
auto result = filter_.filter(std::move(tokens));
ASSERT_EQ(result.size(), 2u);
EXPECT_EQ(result[0].text, "already");
EXPECT_EQ(result[1].text, "lower");
}
TEST_F(LowercaseTokenFilterTest, EmptyToken) {
auto tokens = make_tokens({""});
auto result = filter_.filter(std::move(tokens));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "");
}
TEST_F(LowercaseTokenFilterTest, EmptyList) {
std::vector<Token> tokens;
auto result = filter_.filter(std::move(tokens));
EXPECT_TRUE(result.empty());
}
TEST_F(LowercaseTokenFilterTest, LatinExtended) {
// German uppercase with umlauts
auto tokens =
make_tokens({"\xC3\x9C"
"BER"}); // "ÜBER"
auto result = filter_.filter(std::move(tokens));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text,
"\xC3\xBC"
"ber"); // "über"
}
TEST_F(LowercaseTokenFilterTest, Cyrillic) {
// "МОСКВА" -> "москва"
auto tokens =
make_tokens({"\xD0\x9C\xD0\x9E\xD0\xA1\xD0\x9A\xD0\x92\xD0\x90"});
auto result = filter_.filter(std::move(tokens));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "\xD0\xBC\xD0\xBE\xD1\x81\xD0\xBA\xD0\xB2\xD0\xB0");
}
TEST_F(LowercaseTokenFilterTest, Greek) {
// "ΔΕΛΤΑ" -> "δελτα"
auto tokens = make_tokens({"\xCE\x94\xCE\x95\xCE\x9B\xCE\xA4\xCE\x91"});
auto result = filter_.filter(std::move(tokens));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "\xCE\xB4\xCE\xB5\xCE\xBB\xCF\x84\xCE\xB1");
}
TEST_F(LowercaseTokenFilterTest, MixedScripts) {
// "Hello Мир" -> "hello мир"
auto tokens = make_tokens({"Hello \xD0\x9C\xD0\xB8\xD1\x80"});
auto result = filter_.filter(std::move(tokens));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "hello \xD0\xBC\xD0\xB8\xD1\x80");
}
TEST_F(LowercaseTokenFilterTest, NumbersAndPunctuation) {
auto tokens = make_tokens({"ABC123!@#"});
auto result = filter_.filter(std::move(tokens));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "abc123!@#");
}
TEST_F(LowercaseTokenFilterTest, CJKPassthrough) {
// CJK characters have no case — should pass through unchanged
auto tokens = make_tokens({"\xE4\xB8\xAD\xE6\x96\x87"}); // "中文"
auto result = filter_.filter(std::move(tokens));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "\xE4\xB8\xAD\xE6\x96\x87");
}
TEST_F(LowercaseTokenFilterTest, PreservesOffsetAndPosition) {
std::vector<Token> tokens = {{"ABC", 5, 3}};
auto result = filter_.filter(std::move(tokens));
ASSERT_EQ(result.size(), 1u);
EXPECT_EQ(result[0].text, "abc");
EXPECT_EQ(result[0].offset, 5);
EXPECT_EQ(result[0].position, 3);
}
@@ -0,0 +1,485 @@
// 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 <string>
#include <vector>
#include <gtest/gtest.h>
#include "db/index/column/fts_column/fts_types.h"
#include "db/index/column/fts_column/tokenizer/tokenizer_factory.h"
using namespace zvec::fts;
static std::vector<std::string> token_texts(const std::vector<Token> &tokens) {
std::vector<std::string> texts;
texts.reserve(tokens.size());
for (const auto &token : tokens) {
texts.push_back(token.text);
}
return texts;
}
class StandardTokenizerTest : public ::testing::Test {
protected:
void SetUp() override {
FtsIndexParams params;
params.tokenizer_name = "standard";
params.filters.clear();
pipeline_ = TokenizerFactory::create(params);
ASSERT_NE(pipeline_, nullptr);
}
std::vector<Token> tokenize(const std::string &text) {
return pipeline_->process(text);
}
TokenizerPipelinePtr pipeline_;
};
// --- ASCII basics (existing behavior preserved) ---
TEST_F(StandardTokenizerTest, SimpleAsciiWords) {
auto tokens = tokenize("hello world");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "hello");
EXPECT_EQ(tokens[1].text, "world");
}
TEST_F(StandardTokenizerTest, PunctuationAsDelimiter) {
auto tokens = tokenize("hello,world!test");
ASSERT_EQ(tokens.size(), 3u);
EXPECT_EQ(tokens[0].text, "hello");
EXPECT_EQ(tokens[1].text, "world");
EXPECT_EQ(tokens[2].text, "test");
}
TEST_F(StandardTokenizerTest, LettersAndDigitsTogether) {
auto tokens = tokenize("abc123 xyz");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "abc123");
EXPECT_EQ(tokens[1].text, "xyz");
}
TEST_F(StandardTokenizerTest, EmptyInput) {
auto tokens = tokenize("");
EXPECT_TRUE(tokens.empty());
}
TEST_F(StandardTokenizerTest, OnlyDelimiters) {
auto tokens = tokenize(" .,;! ");
EXPECT_TRUE(tokens.empty());
}
TEST_F(StandardTokenizerTest, MalformedUtf8BreaksTokens) {
std::string text = "ab";
text.push_back(static_cast<char>(0xFF));
text += "cd";
auto tokens = tokenize(text);
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "ab");
EXPECT_EQ(tokens[0].offset, 0u);
EXPECT_EQ(tokens[1].text, "cd");
EXPECT_EQ(tokens[1].offset, 3u);
}
TEST_F(StandardTokenizerTest, OffsetAndPosition) {
auto tokens = tokenize(" hello world");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].offset, 2u);
EXPECT_EQ(tokens[0].position, 0u);
EXPECT_EQ(tokens[1].offset, 8u);
EXPECT_EQ(tokens[1].position, 1u);
}
// --- Accented Latin ---
TEST_F(StandardTokenizerTest, AccentedLatin) {
// café résumé → ["café", "résumé"]
auto tokens = tokenize("caf\xC3\xA9 r\xC3\xA9sum\xC3\xA9");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "caf\xC3\xA9");
EXPECT_EQ(tokens[1].text, "r\xC3\xA9sum\xC3\xA9");
}
TEST_F(StandardTokenizerTest, MarksContinueButDoNotStartTokens) {
// e + U+0301 keeps the combining mark with the base letter.
// Standalone U+0301 and U+FE0F are not indexed.
auto tokens = tokenize(
"e\xCC\x81 "
"\xCC\x81 "
"\xEF\xB8\x8F");
ASSERT_EQ(tokens.size(), 1u);
EXPECT_EQ(tokens[0].text, "e\xCC\x81");
}
TEST_F(StandardTokenizerTest, MaxTokenLengthDoesNotCreateMarkOnlyToken) {
FtsIndexParams params;
params.tokenizer_name = "standard";
params.filters.clear();
params.extra_params = R"({"max_token_length":2})";
auto pipeline = TokenizerFactory::create(params);
ASSERT_NE(pipeline, nullptr);
// ab + U+0301 + c should not split into a standalone combining mark token.
auto tokens = pipeline->process(
"ab\xCC\x81"
"c");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "ab\xCC\x81");
EXPECT_EQ(tokens[1].text, "c");
}
TEST_F(StandardTokenizerTest, GermanUmlaut) {
// Über Straße → ["Über", "Straße"]
auto tokens = tokenize(
"\xC3\x9C"
"ber Stra\xC3\x9F"
"e");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text,
"\xC3\x9C"
"ber");
EXPECT_EQ(tokens[1].text,
"Stra\xC3\x9F"
"e");
}
// --- Cyrillic ---
TEST_F(StandardTokenizerTest, Cyrillic) {
// Москва Россия → ["Москва", "Россия"]
auto tokens = tokenize(
"\xD0\x9C\xD0\xBE\xD1\x81\xD0\xBA\xD0\xB2\xD0\xB0 "
"\xD0\xA0\xD0\xBE\xD1\x81\xD1\x81\xD0\xB8\xD1\x8F");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "\xD0\x9C\xD0\xBE\xD1\x81\xD0\xBA\xD0\xB2\xD0\xB0");
EXPECT_EQ(tokens[1].text, "\xD0\xA0\xD0\xBE\xD1\x81\xD1\x81\xD0\xB8\xD1\x8F");
}
// --- CJK single-character tokenization ---
TEST_F(StandardTokenizerTest, CJKSingleChar) {
// 全文检索 → ["全", "文", "检", "索"]
auto tokens = tokenize("\xE5\x85\xA8\xE6\x96\x87\xE6\xA3\x80\xE7\xB4\xA2");
ASSERT_EQ(tokens.size(), 4u);
EXPECT_EQ(tokens[0].text, "\xE5\x85\xA8"); // 全
EXPECT_EQ(tokens[1].text, "\xE6\x96\x87"); // 文
EXPECT_EQ(tokens[2].text, "\xE6\xA3\x80"); // 检
EXPECT_EQ(tokens[3].text, "\xE7\xB4\xA2"); // 索
}
TEST_F(StandardTokenizerTest, CJKWithSpaces) {
// 你 好 → ["你", "好"]
auto tokens = tokenize("\xE4\xBD\xA0 \xE5\xA5\xBD");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "\xE4\xBD\xA0");
EXPECT_EQ(tokens[1].text, "\xE5\xA5\xBD");
}
TEST_F(StandardTokenizerTest, CJKUnicode17ExtensionBlocks) {
// U+2EBF0 (Extension I), U+31350 (Extension H), U+323B0 (Extension J)
// should each be emitted as an individual CJK token.
auto tokens = tokenize("\xF0\xAE\xAF\xB0\xF0\xB1\x8D\x90\xF0\xB2\x8E\xB0");
ASSERT_EQ(tokens.size(), 3u);
EXPECT_EQ(tokens[0].text, "\xF0\xAE\xAF\xB0");
EXPECT_EQ(tokens[1].text, "\xF0\xB1\x8D\x90");
EXPECT_EQ(tokens[2].text, "\xF0\xB2\x8E\xB0");
}
TEST_F(StandardTokenizerTest, CJKCompatibilitySupplement) {
// U+2F800 CJK Compatibility Ideographs Supplement.
auto tokens = tokenize("\xF0\xAF\xA0\x80");
ASSERT_EQ(tokens.size(), 1u);
EXPECT_EQ(tokens[0].text, "\xF0\xAF\xA0\x80");
}
TEST_F(StandardTokenizerTest, CJKSingleCharKeepsTrailingMarks) {
auto tokens = tokenize("\xE4\xB8\xAD\xEF\xB8\x80\xE6\x96\x87");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "\xE4\xB8\xAD\xEF\xB8\x80");
EXPECT_EQ(tokens[1].text, "\xE6\x96\x87");
}
// --- Mixed scripts ---
TEST_F(StandardTokenizerTest, MixedLatinAndCJK) {
// hello世界test → ["hello", "世", "界", "test"]
auto tokens = tokenize("hello\xE4\xB8\x96\xE7\x95\x8Ctest");
ASSERT_EQ(tokens.size(), 4u);
EXPECT_EQ(tokens[0].text, "hello");
EXPECT_EQ(tokens[1].text, "\xE4\xB8\x96"); // 世
EXPECT_EQ(tokens[2].text, "\xE7\x95\x8C"); // 界
EXPECT_EQ(tokens[3].text, "test");
}
TEST_F(StandardTokenizerTest, CJKWithLatinAndDigits) {
// ES标准分词器v2 → ["ES", "标", "准", "分", "词", "器", "v2"]
auto tokens = tokenize(
"ES\xE6\xA0\x87\xE5\x87\x86\xE5\x88\x86"
"\xE8\xAF\x8D\xE5\x99\xA8v2");
ASSERT_EQ(tokens.size(), 7u);
EXPECT_EQ(tokens[0].text, "ES");
EXPECT_EQ(tokens[1].text, "\xE6\xA0\x87"); // 标
EXPECT_EQ(tokens[2].text, "\xE5\x87\x86"); // 准
EXPECT_EQ(tokens[3].text, "\xE5\x88\x86"); // 分
EXPECT_EQ(tokens[4].text, "\xE8\xAF\x8D"); // 词
EXPECT_EQ(tokens[5].text, "\xE5\x99\xA8"); // 器
EXPECT_EQ(tokens[6].text, "v2");
}
// --- Consecutive positions ---
TEST_F(StandardTokenizerTest, CJKPositionsAreConsecutive) {
auto tokens = tokenize("\xE4\xB8\xAD\xE6\x96\x87"); // 中文
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].position, 0u);
EXPECT_EQ(tokens[1].position, 1u);
}
TEST_F(StandardTokenizerTest, CJKRespectsMaxTokenLength) {
// With max_token_length=1, multi-codepoint words are split.
// CJK chars are always 1 codepoint each — unaffected.
FtsIndexParams params;
params.tokenizer_name = "standard";
params.filters.clear();
params.extra_params = R"({"max_token_length":1})";
auto pipeline = TokenizerFactory::create(params);
ASSERT_NE(pipeline, nullptr);
// "a中bc" → "a", "中", "b", "c" (bc split into b and c)
auto tokens = pipeline->process(
"a\xE4\xB8\xAD"
"bc");
ASSERT_EQ(tokens.size(), 4u);
EXPECT_EQ(tokens[0].text, "a");
EXPECT_EQ(tokens[1].text, "\xE4\xB8\xAD");
EXPECT_EQ(tokens[2].text, "b");
EXPECT_EQ(tokens[3].text, "c");
}
TEST_F(StandardTokenizerTest, MaxTokenLengthSplitsLongWords) {
// "abcdefgh" with max_token_length=5 → ["abcde", "fgh"]
FtsIndexParams params;
params.tokenizer_name = "standard";
params.filters.clear();
params.extra_params = R"({"max_token_length":5})";
auto pipeline = TokenizerFactory::create(params);
ASSERT_NE(pipeline, nullptr);
auto tokens = pipeline->process("abcdefgh");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "abcde");
EXPECT_EQ(tokens[1].text, "fgh");
}
TEST_F(StandardTokenizerTest, MaxTokenLengthCountsCodepointsNotBytes) {
// "café" is 4 codepoints but 5 bytes.
// With max_token_length=4 it fits in one token.
FtsIndexParams params4;
params4.tokenizer_name = "standard";
params4.filters.clear();
params4.extra_params = R"({"max_token_length":4})";
auto pipeline4 = TokenizerFactory::create(params4);
ASSERT_NE(pipeline4, nullptr);
auto tokens4 = pipeline4->process("caf\xC3\xA9");
ASSERT_EQ(tokens4.size(), 1u);
EXPECT_EQ(tokens4[0].text, "caf\xC3\xA9");
// With max_token_length=3 it splits into ["caf", "é"].
FtsIndexParams params3;
params3.tokenizer_name = "standard";
params3.filters.clear();
params3.extra_params = R"({"max_token_length":3})";
auto pipeline3 = TokenizerFactory::create(params3);
ASSERT_NE(pipeline3, nullptr);
auto tokens3 = pipeline3->process("caf\xC3\xA9");
ASSERT_EQ(tokens3.size(), 2u);
EXPECT_EQ(tokens3[0].text, "caf");
EXPECT_EQ(tokens3[1].text, "\xC3\xA9");
}
TEST_F(StandardTokenizerTest, MaxTokenLengthDropsConnectorOnlySplitSegments) {
FtsIndexParams params3;
params3.tokenizer_name = "standard";
params3.filters.clear();
params3.extra_params = R"({"max_token_length":3})";
auto pipeline3 = TokenizerFactory::create(params3);
ASSERT_NE(pipeline3, nullptr);
auto tokens3 = pipeline3->process("dog's");
ASSERT_EQ(tokens3.size(), 2u);
EXPECT_EQ(tokens3[0].text, "dog");
EXPECT_EQ(tokens3[1].text, "s");
FtsIndexParams params1;
params1.tokenizer_name = "standard";
params1.filters.clear();
params1.extra_params = R"({"max_token_length":1})";
auto pipeline1 = TokenizerFactory::create(params1);
ASSERT_NE(pipeline1, nullptr);
auto leading = pipeline1->process("_lead");
std::vector<std::string> expected_leading = {"l", "e", "a", "d"};
EXPECT_EQ(token_texts(leading), expected_leading);
auto internal = pipeline1->process("abc__def");
std::vector<std::string> expected_internal = {"a", "b", "c", "d", "e", "f"};
EXPECT_EQ(token_texts(internal), expected_internal);
}
TEST_F(StandardTokenizerTest, IntraWordPunctuation) {
auto tokens = tokenize(
"dog's 3.14 1,000 example.com hello,world host:port a:b "
"host:9200");
std::vector<std::string> expected = {
"dog's", "3.14", "1,000", "example.com", "hello",
"world", "host:port", "a:b", "host", "9200"};
EXPECT_EQ(token_texts(tokens), expected);
}
TEST_F(StandardTokenizerTest, ExtendNumLetConnectsWordsAndNumbers) {
auto tokens = tokenize("foo_bar v1_2 _lead __123 _");
std::vector<std::string> expected = {"foo_bar", "v1_2", "_lead", "__123"};
EXPECT_EQ(token_texts(tokens), expected);
}
TEST_F(StandardTokenizerTest, EmojiZwjSequence) {
auto tokens = tokenize(
"\xF0\x9F\x91\xA9\xE2\x80\x8D\xF0\x9F\x92\xBB "
"\xE2\x9D\xA4\xEF\xB8\x8F");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "\xF0\x9F\x91\xA9\xE2\x80\x8D\xF0\x9F\x92\xBB");
EXPECT_EQ(tokens[1].text, "\xE2\x9D\xA4\xEF\xB8\x8F");
}
TEST_F(StandardTokenizerTest, EmojiKeycapSequences) {
auto tokens = tokenize(
"1\xEF\xB8\x8F\xE2\x83\xA3 "
"#\xE2\x83\xA3 "
"*\xEF\xB8\x8F\xE2\x83\xA3");
ASSERT_EQ(tokens.size(), 3u);
EXPECT_EQ(tokens[0].text, "1\xEF\xB8\x8F\xE2\x83\xA3");
EXPECT_EQ(tokens[1].text, "#\xE2\x83\xA3");
EXPECT_EQ(tokens[2].text, "*\xEF\xB8\x8F\xE2\x83\xA3");
}
TEST_F(StandardTokenizerTest, EmojiModifierSequences) {
auto tokens = tokenize(
"\xF0\x9F\x91\x8D\xF0\x9F\x8F\xBD "
"\xE2\x98\x9D\xEF\xB8\x8F\xF0\x9F\x8F\xBB "
"\xF0\x9F\x8F\xBD");
ASSERT_EQ(tokens.size(), 3u);
EXPECT_EQ(tokens[0].text, "\xF0\x9F\x91\x8D\xF0\x9F\x8F\xBD");
EXPECT_EQ(tokens[1].text, "\xE2\x98\x9D\xEF\xB8\x8F\xF0\x9F\x8F\xBB");
EXPECT_EQ(tokens[2].text, "\xF0\x9F\x8F\xBD");
}
TEST_F(StandardTokenizerTest, EmojiModifierInsideZwjSequence) {
auto tokens =
tokenize("\xF0\x9F\x91\xA9\xF0\x9F\x8F\xBD\xE2\x80\x8D\xF0\x9F\x92\xBB");
ASSERT_EQ(tokens.size(), 1u);
EXPECT_EQ(tokens[0].text,
"\xF0\x9F\x91\xA9\xF0\x9F\x8F\xBD\xE2\x80\x8D\xF0\x9F\x92\xBB");
}
TEST_F(StandardTokenizerTest, RegionalIndicatorPairs) {
auto tokens = tokenize(
"\xF0\x9F\x87\xBA\xF0\x9F\x87\xB8"
"\xF0\x9F\x87\xA8\xF0\x9F\x87\xA6"
"\xF0\x9F\x87\xAF");
ASSERT_EQ(tokens.size(), 3u);
EXPECT_EQ(tokens[0].text, "\xF0\x9F\x87\xBA\xF0\x9F\x87\xB8");
EXPECT_EQ(tokens[1].text, "\xF0\x9F\x87\xA8\xF0\x9F\x87\xA6");
EXPECT_EQ(tokens[2].text, "\xF0\x9F\x87\xAF");
}
TEST_F(StandardTokenizerTest, RegionalIndicatorPairsIgnoreExtendAndZwj) {
auto tokens = tokenize(
"\xF0\x9F\x87\xA6\xCC\x88\xF0\x9F\x87\xA7 "
"\xF0\x9F\x87\xA6\xE2\x80\x8D\xF0\x9F\x87\xA7\xF0\x9F\x87\xA8");
ASSERT_EQ(tokens.size(), 3u);
EXPECT_EQ(tokens[0].text, "\xF0\x9F\x87\xA6\xCC\x88\xF0\x9F\x87\xA7");
EXPECT_EQ(tokens[1].text, "\xF0\x9F\x87\xA6\xE2\x80\x8D\xF0\x9F\x87\xA7");
EXPECT_EQ(tokens[2].text, "\xF0\x9F\x87\xA8");
}
TEST_F(StandardTokenizerTest, MinimalWb3cZwjExtendedPictographic) {
auto tokens = tokenize(
"\xE2\x80\x8D\xF0\x9F\x9B\x91 "
"a\xE2\x80\x8D\xF0\x9F\x9B\x91 "
"\xE2\x80\x8D\xE2\x93\x82");
ASSERT_EQ(tokens.size(), 3u);
EXPECT_EQ(tokens[0].text, "\xE2\x80\x8D\xF0\x9F\x9B\x91");
EXPECT_EQ(tokens[1].text, "a\xE2\x80\x8D\xF0\x9F\x9B\x91");
EXPECT_EQ(tokens[2].text, "\xE2\x80\x8D\xE2\x93\x82");
}
TEST_F(StandardTokenizerTest, HiraganaTokensAreSingleCharacters) {
auto tokens = tokenize("\xE3\x81\x8B\xE3\x82\x99\xE3\x81\xAA");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "\xE3\x81\x8B\xE3\x82\x99");
EXPECT_EQ(tokens[1].text, "\xE3\x81\xAA");
}
TEST_F(StandardTokenizerTest, JapaneseKoreanAndSoutheastAsianScripts) {
auto tokens = tokenize(
"\xE3\x81\xB2\xE3\x82\x89\xE3\x81\x8C\xE3\x81\xAA "
"\xE3\x82\xAB\xE3\x82\xBF\xE3\x82\xAB\xE3\x83\x8A "
"\xED\x95\x9C\xEA\xB5\xAD "
"\xE0\xB9\x84\xE0\xB8\x97\xE0\xB8\xA2 "
"\xE1\x80\x99\xE1\x80\x94");
ASSERT_EQ(tokens.size(), 8u);
EXPECT_EQ(tokens[0].text, "\xE3\x81\xB2");
EXPECT_EQ(tokens[1].text, "\xE3\x82\x89");
EXPECT_EQ(tokens[2].text, "\xE3\x81\x8C");
EXPECT_EQ(tokens[3].text, "\xE3\x81\xAA");
EXPECT_EQ(tokens[4].text, "\xE3\x82\xAB\xE3\x82\xBF\xE3\x82\xAB\xE3\x83\x8A");
EXPECT_EQ(tokens[5].text, "\xED\x95\x9C\xEA\xB5\xAD");
EXPECT_EQ(tokens[6].text, "\xE0\xB9\x84\xE0\xB8\x97\xE0\xB8\xA2");
EXPECT_EQ(tokens[7].text, "\xE1\x80\x99\xE1\x80\x94");
}
TEST_F(StandardTokenizerTest, SoutheastAsianMarksContinueButDoNotStartTokens) {
auto tokens = tokenize(
"\xE0\xB8\x81\xE0\xB8\xB1 "
"\xE0\xB8\xB1");
ASSERT_EQ(tokens.size(), 1u);
EXPECT_EQ(tokens[0].text, "\xE0\xB8\x81\xE0\xB8\xB1");
}
TEST_F(StandardTokenizerTest, HangulSymbolsOutsideWordClassAreIgnored) {
auto tokens = tokenize("\xE3\x89\xA0 \xED\x95\x9C\xEA\xB5\xAD");
ASSERT_EQ(tokens.size(), 1u);
EXPECT_EQ(tokens[0].text, "\xED\x95\x9C\xEA\xB5\xAD");
}
TEST_F(StandardTokenizerTest, HebrewSingleQuoteStaysWithLetter) {
auto tokens = tokenize("\xD7\x90' \xD7\x90\"\xD7\x91");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "\xD7\x90'");
EXPECT_EQ(tokens[1].text, "\xD7\x90\"\xD7\x91");
}
TEST(StandardTokenizerConfigTest, MaxTokenLengthValidation) {
FtsIndexParams params;
params.tokenizer_name = "standard";
params.filters.clear();
params.extra_params = R"({"max_token_length":0})";
EXPECT_EQ(TokenizerFactory::create(params), nullptr);
params.extra_params = R"({"max_token_length":1048577})";
EXPECT_EQ(TokenizerFactory::create(params), nullptr);
params.extra_params = R"({"max_token_length":1})";
EXPECT_NE(TokenizerFactory::create(params), nullptr);
}
@@ -0,0 +1,167 @@
// 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 <string>
#include <vector>
#include <gtest/gtest.h>
#include "db/index/column/fts_column/fts_types.h"
#include "db/index/column/fts_column/tokenizer/tokenizer_factory.h"
using namespace zvec::fts;
// ============================================================
// Helpers
// ============================================================
static FtsIndexParams make_stemmer_params(
const std::string &lang = "",
const std::vector<std::string> &filters = {"lowercase", "stemmer"}) {
FtsIndexParams params;
params.tokenizer_name = "standard";
params.filters = filters;
if (!lang.empty()) {
params.extra_params = R"({"stemmer_lang":")" + lang + R"("})";
}
return params;
}
// ============================================================
// Pipeline creation
// ============================================================
TEST(StemmerTokenFilterTest, CreatePipelineDefaultEnglish) {
auto pipeline = TokenizerFactory::create(make_stemmer_params());
ASSERT_NE(pipeline, nullptr);
}
TEST(StemmerTokenFilterTest, CreatePipelineExplicitLanguage) {
auto pipeline = TokenizerFactory::create(make_stemmer_params("german"));
ASSERT_NE(pipeline, nullptr);
}
TEST(StemmerTokenFilterTest, CreatePipelineInvalidLanguageFails) {
auto pipeline =
TokenizerFactory::create(make_stemmer_params("nonexistent_lang"));
EXPECT_EQ(pipeline, nullptr);
}
// ============================================================
// English stemming
// ============================================================
TEST(StemmerTokenFilterTest, EnglishStemming) {
auto pipeline = TokenizerFactory::create(make_stemmer_params());
ASSERT_NE(pipeline, nullptr);
auto tokens = pipeline->process("running cats easily connection");
ASSERT_EQ(tokens.size(), 4u);
EXPECT_EQ(tokens[0].text, "run");
EXPECT_EQ(tokens[1].text, "cat");
EXPECT_EQ(tokens[2].text, "easili");
EXPECT_EQ(tokens[3].text, "connect");
}
TEST(StemmerTokenFilterTest, AlreadyStemmedWordsUnchanged) {
auto pipeline = TokenizerFactory::create(make_stemmer_params());
ASSERT_NE(pipeline, nullptr);
auto tokens = pipeline->process("run cat");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].text, "run");
EXPECT_EQ(tokens[1].text, "cat");
}
TEST(StemmerTokenFilterTest, EmptyInput) {
auto pipeline = TokenizerFactory::create(make_stemmer_params());
ASSERT_NE(pipeline, nullptr);
auto tokens = pipeline->process("");
EXPECT_TRUE(tokens.empty());
}
TEST(StemmerTokenFilterTest, PreservesOffsetAndPosition) {
auto pipeline = TokenizerFactory::create(make_stemmer_params());
ASSERT_NE(pipeline, nullptr);
auto tokens = pipeline->process("running dogs");
ASSERT_EQ(tokens.size(), 2u);
EXPECT_EQ(tokens[0].position, 0u);
EXPECT_EQ(tokens[1].position, 1u);
EXPECT_EQ(tokens[0].offset, 0u);
EXPECT_EQ(tokens[1].offset, 8u);
}
// ============================================================
// Lowercase + stemmer chain
// ============================================================
TEST(StemmerTokenFilterTest, LowercaseThenStem) {
auto pipeline = TokenizerFactory::create(make_stemmer_params());
ASSERT_NE(pipeline, nullptr);
auto tokens = pipeline->process("Running Cats EASILY");
ASSERT_EQ(tokens.size(), 3u);
EXPECT_EQ(tokens[0].text, "run");
EXPECT_EQ(tokens[1].text, "cat");
EXPECT_EQ(tokens[2].text, "easili");
}
// ============================================================
// Stemmer-only (no lowercase)
// ============================================================
TEST(StemmerTokenFilterTest, StemmerOnlyNoLowercase) {
auto pipeline =
TokenizerFactory::create(make_stemmer_params("", {"stemmer"}));
ASSERT_NE(pipeline, nullptr);
auto tokens = pipeline->process("running");
ASSERT_EQ(tokens.size(), 1u);
EXPECT_EQ(tokens[0].text, "run");
}
// ============================================================
// Non-English language
// ============================================================
TEST(StemmerTokenFilterTest, GermanStemming) {
auto pipeline = TokenizerFactory::create(make_stemmer_params("german"));
ASSERT_NE(pipeline, nullptr);
auto tokens = pipeline->process("laufen");
ASSERT_EQ(tokens.size(), 1u);
EXPECT_EQ(tokens[0].text, "lauf");
}
// ============================================================
// ISO code as language
// ============================================================
TEST(StemmerTokenFilterTest, LanguageByISOCode) {
auto pipeline = TokenizerFactory::create(make_stemmer_params("en"));
ASSERT_NE(pipeline, nullptr);
auto tokens = pipeline->process("running");
ASSERT_EQ(tokens.size(), 1u);
EXPECT_EQ(tokens[0].text, "run");
}
TEST(StemmerTokenFilterTest, PorterAlgorithm) {
auto pipeline = TokenizerFactory::create(make_stemmer_params("porter"));
ASSERT_NE(pipeline, nullptr);
auto tokens = pipeline->process("running");
ASSERT_EQ(tokens.size(), 1u);
EXPECT_EQ(tokens[0].text, "run");
}
@@ -0,0 +1,271 @@
// 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 "db/index/column/fts_column/tokenizer/tokenizer_pipeline_manager.h"
#include <atomic>
#include <memory>
#include <string>
#include <thread>
#include <vector>
#include <gtest/gtest.h>
#include "db/index/column/fts_column/fts_types.h"
using namespace zvec::fts;
// ============================================================
// Helpers
// ============================================================
static FtsIndexParams make_params(const std::string &tokenizer) {
FtsIndexParams params;
params.tokenizer_name = tokenizer;
return params;
}
// ============================================================
// make_key tests
// ============================================================
TEST(TokenizerPipelineManagerKeyTest, BasicKey) {
FtsIndexParams params;
params.tokenizer_name = "whitespace";
std::string key = TokenizerPipelineManager::make_key(params);
EXPECT_FALSE(key.empty());
EXPECT_NE(key.find("whitespace"), std::string::npos);
}
TEST(TokenizerPipelineManagerKeyTest, SameParamsProduceSameKey) {
FtsIndexParams params1;
params1.tokenizer_name = "whitespace";
params1.extra_params = R"({"dict_path":"/path/to/dict"})";
FtsIndexParams params2;
params2.tokenizer_name = "whitespace";
params2.extra_params = R"({"dict_path":"/path/to/dict"})";
std::string key1 = TokenizerPipelineManager::make_key(params1);
std::string key2 = TokenizerPipelineManager::make_key(params2);
EXPECT_EQ(key1, key2);
}
TEST(TokenizerPipelineManagerKeyTest, DifferentTokenizersDifferentKeys) {
FtsIndexParams params1 = make_params("whitespace");
FtsIndexParams params2 = make_params("jieba");
std::string key1 = TokenizerPipelineManager::make_key(params1);
std::string key2 = TokenizerPipelineManager::make_key(params2);
EXPECT_NE(key1, key2);
}
TEST(TokenizerPipelineManagerKeyTest, FilterNamesAffectKey) {
FtsIndexParams params1 = make_params("whitespace");
params1.filters.clear();
FtsIndexParams params2 = make_params("whitespace");
params2.filters = {"lowercase"};
std::string key1 = TokenizerPipelineManager::make_key(params1);
std::string key2 = TokenizerPipelineManager::make_key(params2);
EXPECT_NE(key1, key2);
}
// ============================================================
// acquire / release tests
// ============================================================
class TokenizerPipelineManagerTest : public ::testing::Test {
protected:
void SetUp() override {
// Use whitespace tokenizer (always available, no dict needed)
params_ = make_params("whitespace");
}
void TearDown() override {
// Best-effort cleanup: release the params if it still exists
// (tests that fail mid-way may leave entries)
// We do this by calling release repeatedly; release on unknown key is a
// no-op
}
FtsIndexParams params_;
};
TEST_F(TokenizerPipelineManagerTest, FirstAcquireCreatesPipeline) {
auto &mgr = TokenizerPipelineManager::Instance();
auto pipeline = mgr.acquire(params_);
ASSERT_NE(pipeline, nullptr);
// Cleanup
mgr.release(params_);
}
TEST_F(TokenizerPipelineManagerTest, RepeatedAcquireReturnsSameInstance) {
auto &mgr = TokenizerPipelineManager::Instance();
auto pipeline1 = mgr.acquire(params_);
auto pipeline2 = mgr.acquire(params_);
ASSERT_NE(pipeline1, nullptr);
ASSERT_NE(pipeline2, nullptr);
// Both should point to the exact same underlying object
EXPECT_EQ(pipeline1.get(), pipeline2.get());
// Cleanup: two acquires → two releases
mgr.release(params_);
mgr.release(params_);
}
TEST_F(TokenizerPipelineManagerTest, ReleaseDecrementsRefCount) {
auto &mgr = TokenizerPipelineManager::Instance();
auto pipeline1 = mgr.acquire(params_);
auto pipeline2 = mgr.acquire(params_);
ASSERT_NE(pipeline1, nullptr);
// Release one reference; pipeline should still be alive (ref_count = 1)
mgr.release(params_);
// Acquire again — should still return the same instance (not recreated)
auto pipeline3 = mgr.acquire(params_);
ASSERT_NE(pipeline3, nullptr);
EXPECT_EQ(pipeline1.get(), pipeline3.get());
// Cleanup: we now have ref_count = 2 (pipeline2 + pipeline3)
mgr.release(params_);
mgr.release(params_);
}
TEST_F(TokenizerPipelineManagerTest, RefCountZeroDestroysEntry) {
auto &mgr = TokenizerPipelineManager::Instance();
auto pipeline1 = mgr.acquire(params_);
ASSERT_NE(pipeline1, nullptr);
void *raw_ptr = pipeline1.get();
// Release the only reference → entry should be removed
mgr.release(params_);
// Acquire again → a new pipeline should be created (possibly different
// address)
auto pipeline2 = mgr.acquire(params_);
ASSERT_NE(pipeline2, nullptr);
// The old shared_ptr (pipeline1) still holds the object alive, so raw_ptr
// is still valid, but the manager has created a fresh entry.
// We can't guarantee same/different address, but we can verify it works.
(void)raw_ptr;
// Cleanup
mgr.release(params_);
}
TEST_F(TokenizerPipelineManagerTest, ReleaseUnknownKeyIsNoOp) {
auto &mgr = TokenizerPipelineManager::Instance();
// Should not crash or assert
FtsIndexParams unknown_params;
unknown_params.tokenizer_name = "nonexistent_tokenizer_name";
EXPECT_NO_THROW(mgr.release(unknown_params));
}
TEST_F(TokenizerPipelineManagerTest, DifferentConfigsDifferentPipelines) {
auto &mgr = TokenizerPipelineManager::Instance();
FtsIndexParams params_ws = make_params("whitespace");
// scws tokenizer will fail to create (no dict), but whitespace should succeed
auto pipeline_ws = mgr.acquire(params_ws);
ASSERT_NE(pipeline_ws, nullptr);
// Cleanup
mgr.release(params_ws);
}
// ============================================================
// Concurrent safety tests
// ============================================================
TEST_F(TokenizerPipelineManagerTest, ConcurrentAcquireSameKey) {
auto &mgr = TokenizerPipelineManager::Instance();
constexpr int kThreads = 8;
constexpr int kAcquiresPerThread = 10;
std::vector<TokenizerPipelinePtr> results(kThreads * kAcquiresPerThread);
std::vector<std::thread> threads;
std::atomic<int> success_count{0};
for (int t = 0; t < kThreads; ++t) {
threads.emplace_back([&, t]() {
for (int i = 0; i < kAcquiresPerThread; ++i) {
auto pipeline = mgr.acquire(params_);
if (pipeline) {
results[t * kAcquiresPerThread + i] = pipeline;
success_count.fetch_add(1);
}
}
});
}
for (auto &th : threads) {
th.join();
}
// All acquires should succeed
EXPECT_EQ(success_count.load(), kThreads * kAcquiresPerThread);
// All non-null results should point to the same underlying pipeline
void *expected_ptr = nullptr;
for (const auto &p : results) {
if (p) {
if (expected_ptr == nullptr) {
expected_ptr = p.get();
} else {
EXPECT_EQ(p.get(), expected_ptr);
}
}
}
// Cleanup: release all acquired references
for (int i = 0; i < kThreads * kAcquiresPerThread; ++i) {
mgr.release(params_);
}
}
TEST_F(TokenizerPipelineManagerTest, ConcurrentAcquireAndRelease) {
auto &mgr = TokenizerPipelineManager::Instance();
constexpr int kThreads = 4;
constexpr int kIterations = 20;
std::atomic<int> errors{0};
std::vector<std::thread> threads;
for (int t = 0; t < kThreads; ++t) {
threads.emplace_back([&]() {
for (int i = 0; i < kIterations; ++i) {
auto pipeline = mgr.acquire(params_);
if (!pipeline) {
errors.fetch_add(1);
continue;
}
// Hold briefly then release
mgr.release(params_);
}
});
}
for (auto &th : threads) {
th.join();
}
EXPECT_EQ(errors.load(), 0);
// After all threads finish, ref_count should be 0 (all released)
// Verify by acquiring once more — should succeed
auto pipeline = mgr.acquire(params_);
EXPECT_NE(pipeline, nullptr);
mgr.release(params_);
}
@@ -0,0 +1,348 @@
// 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 <gtest/gtest.h>
#include "db/index/column/inverted_column/inverted_indexer.h"
#include "tests/test_util.h"
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
#endif
using namespace zvec;
using File = ailego::File;
const std::string working_dir{"./inverted_column_indexer_array_numbers_dir/"};
const std::string collection_name{"test_collection"};
/**
* @brief A helper class for testing the InvertedColumnIndexer implementation.
*
* This class generates test data with specific patterns to verify the
* correctness of the inverted index implementation. It provides various methods
* to populate an InvertedColumnIndexer with predictable data patterns and
* verify that the indexing and search operations work correctly.
*
*/
class TestHelper {
public:
TestHelper(uint32_t num_docs, uint32_t num_write_threads = 10)
: num_docs_(num_docs / 100 * 100),
num_write_threads_(num_write_threads) {};
template <typename T>
void insert_arrays(InvertedColumnIndexer::Ptr indexer) {
auto insert_func = [&](uint32_t start, uint32_t end) {
Status s;
for (uint32_t i = start; i < end; ++i) {
auto arr = generate_array<T>(i);
if (i % 100 == 0) { // Null value for every 100th doc
s = indexer->insert_null(i);
} else {
s = indexer->insert(
i, std::string(reinterpret_cast<const char *>(arr.data()),
sizeof(T) * arr.size()));
}
ASSERT_TRUE(s.ok());
}
};
uint32_t num_docs_per_thread = num_docs_ / num_write_threads_;
std::vector<std::thread> threads{};
for (uint32_t t = 0; t < num_write_threads_; ++t) {
threads.emplace_back(insert_func, t * num_docs_per_thread,
(t + 1) * num_docs_per_thread);
}
for (auto &t : threads) {
t.join();
}
}
template <typename T>
void verify_arrays(InvertedColumnIndexer::Ptr indexer) {
std::vector<std::string> values;
InvertedSearchResult::Ptr res;
// Search for a non-existent value
T v = num_docs_ + 100;
values.emplace_back(std::string((char *)&v, sizeof(T)));
res = indexer->multi_search(values, CompareOp::CONTAIN_ANY);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
// Search for docs containing value "2"
values.clear();
v = 2;
values.emplace_back(std::string((char *)&v, sizeof(T)));
res = indexer->multi_search(values, CompareOp::CONTAIN_ANY);
ASSERT_TRUE(res);
// doc1 and doc2 contain value "2", doc0 is null
ASSERT_EQ(res->count(), 2);
ASSERT_TRUE(res->contains(1));
ASSERT_TRUE(res->contains(2));
res = indexer->multi_search(values, CompareOp::CONTAIN_ALL);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 2);
ASSERT_TRUE(res->contains(1));
ASSERT_TRUE(res->contains(2));
// Search for docs containing values of "2", "3" and "10"
values.clear();
v = 2;
values.emplace_back(std::string((char *)&v, sizeof(T)));
v = 3;
values.emplace_back(std::string((char *)&v, sizeof(T)));
v = 10;
values.emplace_back(std::string((char *)&v, sizeof(T)));
res = indexer->multi_search(values, CompareOp::CONTAIN_ANY);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 8);
ASSERT_TRUE(res->contains(1));
ASSERT_TRUE(res->contains(2));
ASSERT_TRUE(res->contains(3));
ASSERT_TRUE(res->contains(6));
ASSERT_TRUE(res->contains(7));
ASSERT_TRUE(res->contains(8));
ASSERT_TRUE(res->contains(9));
ASSERT_TRUE(res->contains(10));
res = indexer->multi_search(values, CompareOp::CONTAIN_ALL);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
// Search for docs containing values of "3" and "6"
values.clear();
v = 3;
values.emplace_back(std::string((char *)&v, sizeof(T)));
v = 6;
values.emplace_back(std::string((char *)&v, sizeof(T)));
res = indexer->multi_search(values, CompareOp::CONTAIN_ANY);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 6);
ASSERT_TRUE(res->contains(1));
ASSERT_TRUE(res->contains(2));
ASSERT_TRUE(res->contains(3));
ASSERT_TRUE(res->contains(4));
ASSERT_TRUE(res->contains(5));
ASSERT_TRUE(res->contains(6));
res = indexer->multi_search(values, CompareOp::CONTAIN_ALL);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 2);
ASSERT_TRUE(res->contains(2));
ASSERT_TRUE(res->contains(3));
// Search for docs not containing value "1"
values.clear();
v = 1;
values.emplace_back(std::string((char *)&v, sizeof(T)));
res = indexer->multi_search(values, CompareOp::NOT_CONTAIN_ANY);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ - (num_docs_ / 100) - 1);
ASSERT_FALSE(res->contains(1));
res = indexer->multi_search(values, CompareOp::NOT_CONTAIN_ALL);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ - (num_docs_ / 100) - 1);
ASSERT_FALSE(res->contains(1));
// Search for docs not containing value "10" and "14"
values.clear();
v = 10;
values.emplace_back(std::string((char *)&v, sizeof(T)));
v = 14;
values.emplace_back(std::string((char *)&v, sizeof(T)));
res = indexer->multi_search(values, CompareOp::NOT_CONTAIN_ANY);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ - (num_docs_ / 100) - 9);
for (uint32_t id = 6; id <= 14; ++id) {
ASSERT_FALSE(res->contains(id));
}
res = indexer->multi_search(values, CompareOp::NOT_CONTAIN_ALL);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ - (num_docs_ / 100) - 1);
ASSERT_FALSE(res->contains(10));
// Search for docs with array length of 5
res = indexer->search_array_len(5, CompareOp::EQ);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 1000 - (1000 / 100));
res = indexer->search_array_len(5, CompareOp::NE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ - (num_docs_ / 100) - 990);
res = indexer->search_array_len(6, CompareOp::LT);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 1000 - (1000 / 100));
res = indexer->search_array_len(6, CompareOp::LE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ - (num_docs_ / 100));
res = indexer->search_array_len(6, CompareOp::GT);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
res = indexer->search_array_len(6, CompareOp::GE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ - (num_docs_ / 100) - 990);
}
private:
template <typename T>
std::vector<T> generate_array(uint32_t doc_id) {
std::vector<T> nums;
for (uint32_t i = 0; i < 5; ++i) {
T v = doc_id + i;
nums.push_back(v);
}
if (doc_id > 999) {
T v = doc_id + 5;
nums.push_back(v);
}
return nums;
}
private:
const uint32_t num_docs_;
const uint32_t num_write_threads_;
};
/**
*
* @brief Unit tests for the InvertedColumnIndexer implementation.
*
*/
class InvertedIndexTest : public testing::Test {
/***** Global initialization and cleanup - Start *****/
public:
static void SetUpTestCase() {
zvec::test_util::RemoveTestPath(working_dir);
indexer_ = InvertedIndexer::CreateAndOpen(collection_name, working_dir,
true, {}, false);
params_ = std::make_shared<InvertIndexParams>(true);
}
static void TearDownTestCase() {
indexer_.reset();
zvec::test_util::RemoveTestPath(working_dir);
}
/***** Global initialization and cleanup - End *****/
/***** Per-test initialization and cleanup - Start *****/
protected:
void SetUp() override {}
void TearDown() override {}
/***** Per-test initialization and cleanup - End *****/
protected:
static InvertedIndexer::Ptr indexer_;
static TestHelper test_helper_;
static IndexParams::Ptr params_;
};
InvertedIndexer::Ptr InvertedIndexTest::indexer_{nullptr};
TestHelper InvertedIndexTest::test_helper_{100000, 10};
IndexParams::Ptr InvertedIndexTest::params_{nullptr};
/*
*
* Test Cases
*
*/
TEST_F(InvertedIndexTest, ARRAY_INT32) {
ASSERT_TRUE(indexer_);
FieldSchema array_int32{"array_int32", DataType::ARRAY_INT32, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(array_int32).ok());
auto indexer_int32 = (*indexer_)["array_int32"];
ASSERT_TRUE(indexer_int32);
test_helper_.insert_arrays<int32_t>(indexer_int32);
test_helper_.verify_arrays<int32_t>(indexer_int32);
}
TEST_F(InvertedIndexTest, ARRAY_INT64) {
ASSERT_TRUE(indexer_);
FieldSchema array_int64{"array_int64", DataType::ARRAY_INT64, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(array_int64).ok());
auto indexer_int64 = (*indexer_)["array_int64"];
ASSERT_TRUE(indexer_int64);
test_helper_.insert_arrays<int64_t>(indexer_int64);
test_helper_.verify_arrays<int64_t>(indexer_int64);
}
TEST_F(InvertedIndexTest, ARRAY_UINT32) {
ASSERT_TRUE(indexer_);
FieldSchema array_uint32{"array_uint32", DataType::ARRAY_UINT32, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(array_uint32).ok());
auto indexer_uint32 = (*indexer_)["array_uint32"];
ASSERT_TRUE(indexer_uint32);
test_helper_.insert_arrays<uint32_t>(indexer_uint32);
test_helper_.verify_arrays<uint32_t>(indexer_uint32);
}
TEST_F(InvertedIndexTest, ARRAY_UINT64) {
ASSERT_TRUE(indexer_);
FieldSchema array_uint64{"array_uint64", DataType::ARRAY_UINT64, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(array_uint64).ok());
auto indexer_uint64 = (*indexer_)["array_uint64"];
ASSERT_TRUE(indexer_uint64);
test_helper_.insert_arrays<uint64_t>(indexer_uint64);
test_helper_.verify_arrays<uint64_t>(indexer_uint64);
}
TEST_F(InvertedIndexTest, SEALED) {
ASSERT_TRUE(indexer_);
ASSERT_TRUE(indexer_->seal().ok());
auto indexer_int32 = (*indexer_)["array_int32"];
ASSERT_TRUE(indexer_int32);
test_helper_.verify_arrays<int32_t>(indexer_int32);
auto indexer_int64 = (*indexer_)["array_int64"];
ASSERT_TRUE(indexer_int64);
test_helper_.verify_arrays<int64_t>(indexer_int64);
auto indexer_uint32 = (*indexer_)["array_uint32"];
ASSERT_TRUE(indexer_uint32);
test_helper_.verify_arrays<uint32_t>(indexer_uint32);
auto indexer_uint64 = (*indexer_)["array_uint64"];
ASSERT_TRUE(indexer_uint64);
test_helper_.verify_arrays<uint64_t>(indexer_uint64);
}
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif
@@ -0,0 +1,366 @@
// 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 <gtest/gtest.h>
#include "db/index/column/inverted_column/inverted_indexer.h"
#include "tests/test_util.h"
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
#endif
using namespace zvec;
using File = ailego::File;
const std::string working_dir{"./inverted_column_indexer_bool_dir/"};
const std::string collection_name{"test_collection"};
/**
* @brief A helper class for testing the InvertedColumnIndexer implementation.
*
* This class generates test data with specific patterns to verify the
* correctness of the inverted index implementation. It provides various methods
* to populate an InvertedColumnIndexer with predictable data patterns and
* verify that the indexing and search operations work correctly.
*
*/
class TestHelper {
public:
TestHelper(uint32_t num_docs, uint32_t num_write_threads = 10)
: num_docs_(num_docs / 100 * 100),
num_write_threads_(num_write_threads) {};
void insert_bools(InvertedColumnIndexer::Ptr indexer) {
auto insert_func = [&](uint32_t start, uint32_t end) {
Status s;
for (uint32_t i = start; i < end; ++i) {
bool v = generate_bool(i);
s = indexer->insert(i, v);
ASSERT_TRUE(s.ok());
}
};
uint32_t num_docs_per_thread = num_docs_ / num_write_threads_;
std::vector<std::thread> threads{};
for (uint32_t t = 0; t < num_write_threads_; ++t) {
threads.emplace_back(insert_func, t * num_docs_per_thread,
(t + 1) * num_docs_per_thread);
}
for (auto &t : threads) {
t.join();
}
}
void verify_bools(InvertedColumnIndexer::Ptr indexer) {
InvertedSearchResult::Ptr res;
res = indexer->search("true", CompareOp::EQ);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 2);
for (uint32_t i = 0; i < num_docs_; ++i) {
if (i % 2 == 0) {
ASSERT_TRUE(res->contains(i));
} else {
ASSERT_FALSE(res->contains(i));
}
}
res = indexer->search("false", CompareOp::NE);
ASSERT_EQ(res->count(), num_docs_ / 2);
for (uint32_t i = 0; i < num_docs_; ++i) {
if (i % 2 == 0) {
ASSERT_TRUE(res->contains(i));
} else {
ASSERT_FALSE(res->contains(i));
}
}
}
void insert_bool_arrays(InvertedColumnIndexer::Ptr indexer) {
auto insert_func = [&](uint32_t start, uint32_t end) {
Status s;
for (uint32_t i = start; i < end; ++i) {
auto v = generate_bool_array(i);
s = indexer->insert(i, v);
ASSERT_TRUE(s.ok());
}
};
uint32_t num_docs_per_thread = num_docs_ / num_write_threads_;
std::vector<std::thread> threads{};
for (uint32_t t = 0; t < num_write_threads_; ++t) {
threads.emplace_back(insert_func, t * num_docs_per_thread,
(t + 1) * num_docs_per_thread);
}
for (auto &t : threads) {
t.join();
}
}
void verify_bool_arrays(InvertedColumnIndexer::Ptr indexer) {
InvertedSearchResult::Ptr res;
res = indexer->multi_search({"true"}, CompareOp::CONTAIN_ALL);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10 * 8);
for (uint32_t i = 0; i < num_docs_; ++i) {
if (i % 10 == 4 || i % 10 == 7) {
ASSERT_FALSE(res->contains(i));
} else {
ASSERT_TRUE(res->contains(i));
}
}
res = indexer->multi_search({"true"}, CompareOp::CONTAIN_ANY);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10 * 8);
for (uint32_t i = 0; i < num_docs_; ++i) {
if (i % 10 == 4 || i % 10 == 7) {
ASSERT_FALSE(res->contains(i));
} else {
ASSERT_TRUE(res->contains(i));
}
}
res = indexer->multi_search({"true", "false"}, CompareOp::CONTAIN_ALL);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10 * 4);
for (uint32_t i = 0; i < num_docs_; ++i) {
if (i % 10 == 2 || i % 10 == 5 || i % 10 == 8 || i % 10 == 9) {
ASSERT_TRUE(res->contains(i));
} else {
ASSERT_FALSE(res->contains(i));
}
}
res = indexer->multi_search({"true", "false"}, CompareOp::CONTAIN_ANY);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_);
res = indexer->search_array_len(1, CompareOp::EQ);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10);
res = indexer->search_array_len(2, CompareOp::EQ);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10 * 2);
res = indexer->search_array_len(3, CompareOp::EQ);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10 * 3);
res = indexer->search_array_len(4, CompareOp::EQ);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10 * 4);
res = indexer->search_array_len(5, CompareOp::NE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_);
res = indexer->search_array_len(3, CompareOp::NE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10 * 7);
res = indexer->search_array_len(1, CompareOp::LT);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
res = indexer->search_array_len(1, CompareOp::LE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10);
res = indexer->search_array_len(4, CompareOp::LT);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10 * 6);
res = indexer->search_array_len(4, CompareOp::LE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_);
res = indexer->search_array_len(1, CompareOp::GT);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10 * 9);
res = indexer->search_array_len(1, CompareOp::GE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_);
res = indexer->search_array_len(4, CompareOp::GT);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
res = indexer->search_array_len(4, CompareOp::GE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10 * 4);
}
private:
bool generate_bool(uint32_t doc_id) {
if (doc_id % 2 == 0) {
return true;
} else {
return false;
}
}
std::vector<bool> generate_bool_array(uint32_t doc_id) {
switch (doc_id % 10) {
case 0:
return {true};
case 1:
return {true, true};
case 2:
return {true, false};
case 3:
return {true, true, true};
case 4:
return {false, false, false};
case 5:
return {false, true, false};
case 6:
return {true, true, true, true};
case 7:
return {false, false, false, false};
case 8:
return {true, false, true, false};
case 9:
return {false, true, false, true};
default:
return {};
}
}
private:
const uint32_t num_docs_;
const uint32_t num_write_threads_;
};
/**
*
* @brief Unit tests for the InvertedColumnIndexer implementation.
*
*/
class InvertedIndexTest : public testing::Test {
/***** Global initialization and cleanup - Start *****/
public:
static void SetUpTestCase() {
zvec::test_util::RemoveTestPath(working_dir);
indexer_ = InvertedIndexer::CreateAndOpen(collection_name, working_dir,
true, {}, false);
params_ = std::make_shared<InvertIndexParams>(true);
}
static void TearDownTestCase() {
indexer_.reset();
zvec::test_util::RemoveTestPath(working_dir);
}
/***** Global initialization and cleanup - End *****/
/***** Per-test initialization and cleanup - Start *****/
protected:
void SetUp() override {}
void TearDown() override {}
/***** Per-test initialization and cleanup - End *****/
protected:
static InvertedIndexer::Ptr indexer_;
static TestHelper test_helper_;
static IndexParams::Ptr params_;
};
InvertedIndexer::Ptr InvertedIndexTest::indexer_{nullptr};
TestHelper InvertedIndexTest::test_helper_{100000, 10};
IndexParams::Ptr InvertedIndexTest::params_{nullptr};
/*
*
* Test Cases
*
*/
TEST_F(InvertedIndexTest, BOOLS) {
ASSERT_TRUE(indexer_);
FieldSchema test_bool{"test_bool", DataType::BOOL, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(test_bool).ok());
auto indexer_bool = (*indexer_)["test_bool"];
ASSERT_TRUE(indexer_bool);
test_helper_.insert_bools(indexer_bool);
test_helper_.verify_bools(indexer_bool);
}
TEST_F(InvertedIndexTest, BOOL_ARRAYS) {
ASSERT_TRUE(indexer_);
FieldSchema test_bool_array{"test_bool_array", DataType::ARRAY_BOOL, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(test_bool_array).ok());
auto indexer_bool_array = (*indexer_)["test_bool_array"];
ASSERT_TRUE(indexer_bool_array);
test_helper_.insert_bool_arrays(indexer_bool_array);
test_helper_.verify_bool_arrays(indexer_bool_array);
}
TEST_F(InvertedIndexTest, SEALED) {
ASSERT_TRUE(indexer_);
ASSERT_TRUE(indexer_->seal().ok());
auto indexer_bool = (*indexer_)["test_bool"];
ASSERT_TRUE(indexer_bool);
test_helper_.verify_bools(indexer_bool);
auto indexer_bool_array = (*indexer_)["test_bool_array"];
ASSERT_TRUE(indexer_bool_array);
test_helper_.verify_bool_arrays(indexer_bool_array);
}
TEST_F(InvertedIndexTest, SNAPSHOT) {
#ifdef __ANDROID__
GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink support (needed by RocksDB checkpoint)";
#endif
ASSERT_TRUE(indexer_);
ASSERT_TRUE(indexer_->create_snapshot(working_dir + "snapshot").ok());
FieldSchema test_bool{"test_bool", DataType::BOOL, true, params_};
FieldSchema test_bool_array{"test_bool_array", DataType::ARRAY_BOOL, true,
params_};
auto snapshot_indexer =
InvertedIndexer::CreateAndOpen(collection_name, working_dir + "snapshot",
false, {test_bool, test_bool_array}, true);
ASSERT_TRUE(snapshot_indexer);
auto indexer_bool = (*snapshot_indexer)["test_bool"];
ASSERT_TRUE(indexer_bool);
test_helper_.verify_bools(indexer_bool);
auto indexer_bool_array = (*snapshot_indexer)["test_bool_array"];
ASSERT_TRUE(indexer_bool_array);
test_helper_.verify_bool_arrays(indexer_bool_array);
}
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif
@@ -0,0 +1,589 @@
// 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 <random>
#include <gtest/gtest.h>
#include "db/index/column/inverted_column/inverted_indexer.h"
#include "tests/test_util.h"
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
#endif
using namespace zvec;
using File = ailego::File;
const std::string working_dir{"./inverted_column_indexer_cyclic_numbers_dir/"};
const std::string collection_name{"test_collection"};
/**
* @brief A helper class for testing the InvertedColumnIndexer implementation.
*
* This class generates test data with specific patterns to verify the
* correctness of the inverted index implementation. It provides various methods
* to populate an InvertedColumnIndexer with predictable data patterns and
* verify that the indexing and search operations work correctly.
*
*/
class TestHelper {
public:
TestHelper(uint32_t num_docs, uint32_t num_write_threads = 10)
: num_docs_(num_docs / 100 * 100),
num_write_threads_(num_write_threads) {};
template <typename T>
void insert_cyclic_numbers(InvertedColumnIndexer::Ptr indexer,
bool include_nulls) {
auto insert_func = [&](uint32_t start, uint32_t end) {
Status s;
for (uint32_t i = start; i < end; ++i) {
T v = generate_cyclic_number<T>(i);
if (include_nulls && i % 100 == 0) { // Null value for every 100th doc
s = indexer->insert_null(i);
} else {
s = indexer->insert(i, std::string((char *)&v, sizeof(T)));
}
ASSERT_TRUE(s.ok());
}
};
uint32_t num_docs_per_thread = num_docs_ / num_write_threads_;
std::vector<std::thread> threads{};
for (uint32_t t = 0; t < num_write_threads_; ++t) {
threads.emplace_back(insert_func, t * num_docs_per_thread,
(t + 1) * num_docs_per_thread);
}
for (auto &t : threads) {
t.join();
}
}
template <typename T>
void verify_cyclic_numbers(InvertedColumnIndexer::Ptr indexer,
bool include_nulls) {
verify_cyclic_numbers_eq_ne<T>(indexer, include_nulls);
verify_cyclic_numbers_range<T>(indexer, include_nulls);
if (include_nulls) {
verify_cyclic_numbers_null<T>(indexer);
}
}
template <typename T>
void verify_cyclic_numbers_eq_ne(InvertedColumnIndexer::Ptr indexer,
bool include_nulls) {
InvertedSearchResult::Ptr res;
// Test EQ operator
for (uint32_t i = 0; i < num_docs_ / 100; ++i) {
uint32_t first_doc_in_cycle = i * 100;
// Search for the first value in this 100-doc cycle
T v = generate_cyclic_number<T>(first_doc_in_cycle);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::EQ);
ASSERT_TRUE(res);
if (include_nulls) {
ASSERT_EQ(res->count(), 9);
for (uint32_t j = 1; j < 10; ++j) {
ASSERT_TRUE(res->contains(first_doc_in_cycle + j * 10));
}
} else {
ASSERT_EQ(res->count(), 10);
for (uint32_t j = 0; j < 10; ++j) {
ASSERT_TRUE(res->contains(first_doc_in_cycle + j * 10));
}
}
// Search for the 4th value in this 100-doc cycle
v = generate_cyclic_number<T>(first_doc_in_cycle + 3);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::EQ);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 10);
for (uint32_t j = 0; j < 10; ++j) {
ASSERT_TRUE(res->contains(first_doc_in_cycle + 3 + j * 10));
}
// Search for an non-existent value
v = first_doc_in_cycle + 11;
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::EQ);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
}
// Test NE operator with a non-existent value
T v = generate_cyclic_number<T>(num_docs_);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::NE);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
// Test NE operator with a random value
static std::random_device rd;
static std::mt19937 gen(rd());
std::uniform_int_distribution<uint32_t> dis(0, num_docs_ / 100 - 1);
uint32_t random_cycle = dis(gen);
v = generate_cyclic_number<T>(random_cycle * 100 + 1);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::NE);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < random_cycle * 100; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
for (uint32_t id = random_cycle * 100; id < (random_cycle + 1) * 100;
++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else if (id % 10 == 1) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
for (uint32_t id = (random_cycle + 1) * 100; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
}
template <typename T>
void verify_cyclic_numbers_range(InvertedColumnIndexer::Ptr indexer,
bool include_nulls) {
InvertedSearchResult::Ptr res;
T v = generate_cyclic_number<T>(0);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LT);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LE);
ASSERT_TRUE(res);
if (include_nulls) {
ASSERT_EQ(res->count(), 9);
} else {
ASSERT_EQ(res->count(), 10);
}
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GT);
ASSERT_TRUE(res);
if (include_nulls) {
ASSERT_EQ(res->count(), num_docs_ - (num_docs_ / 100) - 9);
} else {
ASSERT_EQ(res->count(), num_docs_ - 10);
}
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GE);
ASSERT_TRUE(res);
if (include_nulls) {
ASSERT_EQ(res->count(), num_docs_ - (num_docs_ / 100));
} else {
ASSERT_EQ(res->count(), num_docs_);
}
uint32_t middle_cycle = num_docs_ / 100 / 2;
v = generate_cyclic_number<T>(middle_cycle * 100 + 1);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LT);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < middle_cycle * 100; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
for (uint32_t id = middle_cycle * 100; id < (middle_cycle + 1) * 100;
++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else if (id % 10 < 1) {
ASSERT_TRUE(res->contains(id));
} else {
ASSERT_FALSE(res->contains(id));
}
}
for (uint32_t id = (middle_cycle + 1) * 100; id < num_docs_; ++id) {
ASSERT_FALSE(res->contains(id));
}
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LE);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < middle_cycle * 100; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
for (uint32_t id = middle_cycle * 100; id < (middle_cycle + 1) * 100;
++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else if (id % 10 <= 1) {
ASSERT_TRUE(res->contains(id));
} else {
ASSERT_FALSE(res->contains(id));
}
}
for (uint32_t id = (middle_cycle + 1) * 100; id < num_docs_; ++id) {
ASSERT_FALSE(res->contains(id));
}
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GT);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < middle_cycle * 100; ++id) {
ASSERT_FALSE(res->contains(id));
}
for (uint32_t id = middle_cycle * 100; id < (middle_cycle + 1) * 100;
++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else if (id % 10 > 1) {
ASSERT_TRUE(res->contains(id));
} else {
ASSERT_FALSE(res->contains(id));
}
}
for (uint32_t id = (middle_cycle + 1) * 100; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GE);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < middle_cycle * 100; ++id) {
ASSERT_FALSE(res->contains(id));
}
for (uint32_t id = middle_cycle * 100; id < (middle_cycle + 1) * 100;
++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else if (id % 10 >= 1) {
ASSERT_TRUE(res->contains(id));
} else {
ASSERT_FALSE(res->contains(id));
}
}
for (uint32_t id = (middle_cycle + 1) * 100; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
v = generate_cyclic_number<T>(num_docs_ - 1);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LT);
ASSERT_TRUE(res);
if (include_nulls) {
ASSERT_EQ(res->count(), num_docs_ - (num_docs_ / 100) - 10);
} else {
ASSERT_EQ(res->count(), num_docs_ - 10);
}
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LE);
ASSERT_TRUE(res);
if (include_nulls) {
ASSERT_EQ(res->count(), num_docs_ - (num_docs_ / 100));
} else {
ASSERT_EQ(res->count(), num_docs_);
}
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GT);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 10);
}
template <typename T>
void verify_cyclic_numbers_null(InvertedColumnIndexer::Ptr indexer) {
InvertedSearchResult::Ptr res = indexer->search_null();
ASSERT_TRUE(res);
for (uint32_t i = 0; i < num_docs_; ++i) {
if (i % 100 == 0) {
ASSERT_TRUE(res->contains(i));
} else {
ASSERT_FALSE(res->contains(i));
}
}
res = indexer->search_non_null();
ASSERT_TRUE(res);
for (uint32_t i = 0; i < num_docs_; ++i) {
if (i % 100 == 0) {
ASSERT_FALSE(res->contains(i));
} else {
ASSERT_TRUE(res->contains(i));
}
}
}
private:
template <typename T>
T generate_cyclic_number(uint32_t doc_id) {
// Creates a pattern where every 100 consecutive document IDs share a cycle
// of 10 distinct values.
// E.g., for int32_t,[id: 304, value: 304], [id: 315, value: 305];
// for float, [id: 101, value: 101.666], [id: 112, value: 102.666]
double num_double = (uint32_t)(doc_id / 100) * 100 + doc_id % 10 + 0.666;
T num = num_double;
return num;
}
private:
const uint32_t num_docs_;
const uint32_t num_write_threads_;
};
/**
*
* @brief Unit tests for the InvertedColumnIndexer implementation.
*
*/
class InvertedIndexTest : public testing::Test {
/***** Global initialization and cleanup - Start *****/
public:
static void SetUpTestCase() {
zvec::test_util::RemoveTestPath(working_dir);
indexer_ = InvertedIndexer::CreateAndOpen(collection_name, working_dir,
true, {}, false);
params_ = std::make_shared<InvertIndexParams>(true);
}
static void TearDownTestCase() {
indexer_.reset();
zvec::test_util::RemoveTestPath(working_dir);
}
/***** Global initialization and cleanup - End *****/
/***** Per-test initialization and cleanup - Start *****/
protected:
void SetUp() override {}
void TearDown() override {}
/***** Per-test initialization and cleanup - End *****/
protected:
static InvertedIndexer::Ptr indexer_;
static TestHelper test_helper_;
static IndexParams::Ptr params_;
};
InvertedIndexer::Ptr InvertedIndexTest::indexer_{nullptr};
TestHelper InvertedIndexTest::test_helper_{100000, 10};
IndexParams::Ptr InvertedIndexTest::params_{nullptr};
/*
*
* Test Cases
*
*/
TEST_F(InvertedIndexTest, CYCLIC_NUMBERS_INT32) {
ASSERT_TRUE(indexer_);
FieldSchema cyclic_int32{"cyclic_int32", DataType::INT32, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(cyclic_int32).ok());
auto indexer_int32 = (*indexer_)["cyclic_int32"];
ASSERT_TRUE(indexer_int32);
test_helper_.insert_cyclic_numbers<int32_t>(indexer_int32, false);
test_helper_.verify_cyclic_numbers<int32_t>(indexer_int32, false);
FieldSchema cyclic_int32_w_null{"cyclic_int32_w_null", DataType::INT32, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(cyclic_int32_w_null).ok());
auto indexer_int32_w_null = (*indexer_)["cyclic_int32_w_null"];
ASSERT_TRUE(indexer_int32_w_null);
test_helper_.insert_cyclic_numbers<int32_t>(indexer_int32_w_null, true);
test_helper_.verify_cyclic_numbers<int32_t>(indexer_int32_w_null, true);
}
TEST_F(InvertedIndexTest, CYCLIC_NUMBERS_INT64) {
ASSERT_TRUE(indexer_);
FieldSchema cyclic_int64{"cyclic_int64", DataType::INT64, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(cyclic_int64).ok());
auto indexer_int64 = (*indexer_)["cyclic_int64"];
ASSERT_TRUE(indexer_int64);
test_helper_.insert_cyclic_numbers<int64_t>(indexer_int64, false);
test_helper_.verify_cyclic_numbers<int64_t>(indexer_int64, false);
FieldSchema cyclic_int64_w_null{"cyclic_int64_w_null", DataType::INT64, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(cyclic_int64_w_null).ok());
auto indexer_int64_w_null = (*indexer_)["cyclic_int64_w_null"];
ASSERT_TRUE(indexer_int64_w_null);
test_helper_.insert_cyclic_numbers<int64_t>(indexer_int64_w_null, true);
test_helper_.verify_cyclic_numbers<int64_t>(indexer_int64_w_null, true);
}
TEST_F(InvertedIndexTest, CYCLIC_NUMBERS_UINT32) {
ASSERT_TRUE(indexer_);
FieldSchema cyclic_uint32{"cyclic_uint32", DataType::UINT32, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(cyclic_uint32).ok());
auto indexer_uint32 = (*indexer_)["cyclic_uint32"];
ASSERT_TRUE(indexer_uint32);
test_helper_.insert_cyclic_numbers<uint32_t>(indexer_uint32, false);
test_helper_.verify_cyclic_numbers<uint32_t>(indexer_uint32, false);
FieldSchema cyclic_uint32_w_null{"cyclic_uint32_w_null", DataType::UINT32,
true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(cyclic_uint32_w_null).ok());
auto indexer_uint32_w_null = (*indexer_)["cyclic_uint32_w_null"];
ASSERT_TRUE(indexer_uint32_w_null);
test_helper_.insert_cyclic_numbers<uint32_t>(indexer_uint32_w_null, true);
test_helper_.verify_cyclic_numbers<uint32_t>(indexer_uint32_w_null, true);
}
TEST_F(InvertedIndexTest, CYCLIC_NUMBERS_UINT64) {
ASSERT_TRUE(indexer_);
FieldSchema cyclic_uint64{"cyclic_uint64", DataType::UINT64, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(cyclic_uint64).ok());
auto indexer_uint64 = (*indexer_)["cyclic_uint64"];
ASSERT_TRUE(indexer_uint64);
test_helper_.insert_cyclic_numbers<uint64_t>(indexer_uint64, false);
test_helper_.verify_cyclic_numbers<uint64_t>(indexer_uint64, false);
FieldSchema cyclic_uint64_w_null{"cyclic_uint64_w_null", DataType::UINT64,
true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(cyclic_uint64_w_null).ok());
auto indexer_uint64_w_null = (*indexer_)["cyclic_uint64_w_null"];
ASSERT_TRUE(indexer_uint64_w_null);
test_helper_.insert_cyclic_numbers<uint64_t>(indexer_uint64_w_null, true);
test_helper_.verify_cyclic_numbers<uint64_t>(indexer_uint64_w_null, true);
}
TEST_F(InvertedIndexTest, CYCLIC_NUMBERS_FLOAT) {
ASSERT_TRUE(indexer_);
FieldSchema cyclic_float{"cyclic_float", DataType::FLOAT, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(cyclic_float).ok());
auto indexer_float = (*indexer_)["cyclic_float"];
ASSERT_TRUE(indexer_float);
test_helper_.insert_cyclic_numbers<float>(indexer_float, false);
test_helper_.verify_cyclic_numbers<float>(indexer_float, false);
FieldSchema cyclic_float_w_null{"cyclic_float_w_null", DataType::FLOAT, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(cyclic_float_w_null).ok());
auto indexer_float_w_null = (*indexer_)["cyclic_float_w_null"];
ASSERT_TRUE(indexer_float_w_null);
test_helper_.insert_cyclic_numbers<float>(indexer_float_w_null, true);
test_helper_.verify_cyclic_numbers<float>(indexer_float_w_null, true);
}
TEST_F(InvertedIndexTest, CYCLIC_NUMBERS_DOUBLE) {
ASSERT_TRUE(indexer_);
FieldSchema cyclic_double{"cyclic_double", DataType::DOUBLE, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(cyclic_double).ok());
auto indexer_double = (*indexer_)["cyclic_double"];
ASSERT_TRUE(indexer_double);
test_helper_.insert_cyclic_numbers<double>(indexer_double, false);
test_helper_.verify_cyclic_numbers<double>(indexer_double, false);
FieldSchema cyclic_double_w_null{"cyclic_double_w_null", DataType::DOUBLE,
true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(cyclic_double_w_null).ok());
auto indexer_double_w_null = (*indexer_)["cyclic_double_w_null"];
ASSERT_TRUE(indexer_double_w_null);
test_helper_.insert_cyclic_numbers<double>(indexer_double_w_null, true);
test_helper_.verify_cyclic_numbers<double>(indexer_double_w_null, true);
}
TEST_F(InvertedIndexTest, SEALED) {
ASSERT_TRUE(indexer_);
ASSERT_TRUE(indexer_->seal().ok());
auto indexer_int32 = (*indexer_)["cyclic_int32"];
ASSERT_TRUE(indexer_int32);
test_helper_.verify_cyclic_numbers<int32_t>(indexer_int32, false);
auto indexer_int32_w_null = (*indexer_)["cyclic_int32_w_null"];
ASSERT_TRUE(indexer_int32_w_null);
test_helper_.verify_cyclic_numbers<int32_t>(indexer_int32_w_null, true);
auto indexer_int64 = (*indexer_)["cyclic_int64"];
ASSERT_TRUE(indexer_int64);
test_helper_.verify_cyclic_numbers<int64_t>(indexer_int64, false);
auto indexer_int64_w_null = (*indexer_)["cyclic_int64_w_null"];
ASSERT_TRUE(indexer_int64_w_null);
test_helper_.verify_cyclic_numbers<int64_t>(indexer_int64_w_null, true);
auto indexer_uint32 = (*indexer_)["cyclic_uint32"];
ASSERT_TRUE(indexer_uint32);
test_helper_.verify_cyclic_numbers<uint32_t>(indexer_uint32, false);
auto indexer_uint32_w_null = (*indexer_)["cyclic_uint32_w_null"];
ASSERT_TRUE(indexer_uint32_w_null);
test_helper_.verify_cyclic_numbers<uint32_t>(indexer_uint32_w_null, true);
auto indexer_uint64 = (*indexer_)["cyclic_uint64"];
ASSERT_TRUE(indexer_uint64);
test_helper_.verify_cyclic_numbers<uint64_t>(indexer_uint64, false);
auto indexer_uint64_w_null = (*indexer_)["cyclic_uint64_w_null"];
ASSERT_TRUE(indexer_uint64_w_null);
test_helper_.verify_cyclic_numbers<uint64_t>(indexer_uint64_w_null, true);
auto indexer_float = (*indexer_)["cyclic_float"];
ASSERT_TRUE(indexer_float);
test_helper_.verify_cyclic_numbers<float>(indexer_float, false);
auto indexer_float_w_null = (*indexer_)["cyclic_float_w_null"];
ASSERT_TRUE(indexer_float_w_null);
test_helper_.verify_cyclic_numbers<float>(indexer_float_w_null, true);
auto indexer_double = (*indexer_)["cyclic_double"];
ASSERT_TRUE(indexer_double);
test_helper_.verify_cyclic_numbers<double>(indexer_double, false);
auto indexer_double_w_null = (*indexer_)["cyclic_double_w_null"];
ASSERT_TRUE(indexer_double_w_null);
test_helper_.verify_cyclic_numbers<double>(indexer_double_w_null, true);
}
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif
@@ -0,0 +1,777 @@
// 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 <random>
#include <gtest/gtest.h>
#include "db/index/column/inverted_column/inverted_indexer.h"
#include "tests/test_util.h"
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
#endif
using namespace zvec;
using File = ailego::File;
const std::string working_dir{
"./inverted_column_indexer_sequential_numbers_dir/"};
const std::string collection_name{"test_collection"};
/**
* @brief A helper class for testing the InvertedColumnIndexer implementation.
*
* This class generates test data with specific patterns to verify the
* correctness of the inverted index implementation. It provides various methods
* to populate an InvertedColumnIndexer with predictable data patterns and
* verify that the indexing and search operations work correctly.
*
*/
class TestHelper {
public:
TestHelper(uint32_t num_docs, uint32_t num_write_threads = 10)
: num_docs_(num_docs / 100 * 100),
num_write_threads_(num_write_threads) {};
template <typename T>
void insert_sequential_numbers(InvertedColumnIndexer::Ptr indexer,
bool include_nulls) {
auto insert_func = [&](uint32_t start, uint32_t end) {
Status s;
for (uint32_t i = start; i < end; ++i) {
T v = generate_sequential_number<T>(i);
if (include_nulls && i % 100 == 0) { // Null value for every 100th doc
s = indexer->insert_null(i);
} else {
s = indexer->insert(i, std::string((char *)&v, sizeof(T)));
}
ASSERT_TRUE(s.ok());
}
};
uint32_t num_docs_per_thread = num_docs_ / num_write_threads_;
std::vector<std::thread> threads{};
for (uint32_t t = 0; t < num_write_threads_; ++t) {
threads.emplace_back(insert_func, t * num_docs_per_thread,
(t + 1) * num_docs_per_thread);
}
for (auto &t : threads) {
t.join();
}
}
template <typename T>
void verify_sequential_numbers(InvertedColumnIndexer::Ptr indexer,
bool include_nulls) {
verify_sequential_numbers_eq_ne<T>(indexer, include_nulls);
verify_sequential_numbers_range_less<T>(indexer, include_nulls);
verify_sequential_numbers_range_greater<T>(indexer, include_nulls);
if (include_nulls) {
verify_sequential_numbers_null<T>(indexer);
}
if (indexer->is_sealed()) {
verify_sequential_numbers_range_ratio<T>(indexer, include_nulls);
}
}
template <typename T>
void verify_sequential_numbers_eq_ne(InvertedColumnIndexer::Ptr indexer,
bool include_nulls) {
InvertedSearchResult::Ptr res;
// Test EQ operator
for (uint32_t id = 0; id < num_docs_; ++id) {
T v = generate_sequential_number<T>(id);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::EQ);
ASSERT_TRUE(res);
if (include_nulls && id % 100 == 0) {
ASSERT_EQ(res->count(), 0);
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_EQ(res->count(), 1);
ASSERT_TRUE(res->contains(id));
auto it = res->create_iterator();
ASSERT_EQ(it->doc_id(), id);
it->next();
ASSERT_FALSE(it->valid());
}
}
// Test NE operator with a non-existent value
T v = generate_sequential_number<T>(num_docs_);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::NE);
ASSERT_TRUE(res);
if (include_nulls) {
for (uint32_t id = 0; id < num_docs_; ++id) {
if (id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
} else {
ASSERT_EQ(res->count(), num_docs_);
auto it = res->create_iterator();
for (uint32_t id = 0; id < num_docs_; ++id) {
ASSERT_TRUE(res->contains(id));
ASSERT_EQ(it->doc_id(), id);
it->next();
}
ASSERT_FALSE(it->valid());
}
// Test NE operator with a random value
static std::random_device rd;
static std::mt19937 gen(rd());
std::uniform_int_distribution<uint32_t> dis(0, num_docs_ - 1);
uint32_t num_random = dis(gen);
v = generate_sequential_number<T>(num_random);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::NE);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else if (id == num_random) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
}
template <typename T>
void verify_sequential_numbers_range_less(InvertedColumnIndexer::Ptr indexer,
bool include_nulls) {
T v = generate_sequential_number<T>(0);
auto res =
indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LT);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
ASSERT_FALSE(res->contains(0));
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LE);
ASSERT_TRUE(res);
if (include_nulls) {
ASSERT_EQ(res->count(), 0);
ASSERT_FALSE(res->contains(0));
} else {
ASSERT_EQ(res->count(), 1);
ASSERT_TRUE(res->contains(0));
ASSERT_FALSE(res->contains(1));
}
v = generate_sequential_number<T>(1);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LT);
ASSERT_TRUE(res);
if (include_nulls) {
ASSERT_EQ(res->count(), 0);
ASSERT_FALSE(res->contains(0));
} else {
ASSERT_EQ(res->count(), 1);
ASSERT_TRUE(res->contains(0));
ASSERT_FALSE(res->contains(1));
}
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LE);
ASSERT_TRUE(res);
if (include_nulls) {
ASSERT_EQ(res->count(), 1);
ASSERT_FALSE(res->contains(0));
ASSERT_TRUE(res->contains(1));
ASSERT_FALSE(res->contains(2));
} else {
ASSERT_EQ(res->count(), 2);
ASSERT_TRUE(res->contains(0));
ASSERT_TRUE(res->contains(1));
ASSERT_FALSE(res->contains(2));
}
v = generate_sequential_number<T>(num_docs_ / 10);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LT);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < num_docs_ / 10; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
ASSERT_FALSE(res->contains(num_docs_ / 10));
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LE);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < num_docs_ / 10 + 1; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
ASSERT_FALSE(res->contains(num_docs_ / 10 + 1));
v = generate_sequential_number<T>(num_docs_ / 2);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LT);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < num_docs_ / 2; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
ASSERT_FALSE(res->contains(num_docs_ / 2));
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LE);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < num_docs_ / 2 + 1; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
ASSERT_FALSE(res->contains(num_docs_ / 2 + 1));
v = generate_sequential_number<T>(num_docs_ - 1);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LT);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < num_docs_ - 1; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
ASSERT_FALSE(res->contains(num_docs_ - 1));
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LE);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
ASSERT_FALSE(res->contains(num_docs_));
v = generate_sequential_number<T>(num_docs_);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LT);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
ASSERT_FALSE(res->contains(num_docs_));
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::LE);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
ASSERT_FALSE(res->contains(num_docs_));
}
template <typename T>
void verify_sequential_numbers_range_greater(
InvertedColumnIndexer::Ptr indexer, bool include_nulls) {
T v = generate_sequential_number<T>(0);
auto res =
indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GT);
ASSERT_TRUE(res);
ASSERT_FALSE(res->contains(0));
for (uint32_t id = 1; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GE);
ASSERT_TRUE(res);
for (uint32_t id = 0; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
v = generate_sequential_number<T>(1);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GT);
ASSERT_TRUE(res);
ASSERT_FALSE(res->contains(0));
ASSERT_FALSE(res->contains(1));
for (uint32_t id = 2; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GE);
ASSERT_TRUE(res);
ASSERT_FALSE(res->contains(0));
for (uint32_t id = 1; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
v = generate_sequential_number<T>(num_docs_ / 10);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GT);
ASSERT_TRUE(res);
for (uint32_t id = num_docs_ / 10 + 1; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
ASSERT_FALSE(res->contains(num_docs_ / 10));
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GE);
ASSERT_TRUE(res);
for (uint32_t id = num_docs_ / 10; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
ASSERT_FALSE(res->contains(num_docs_ / 10 - 1));
v = generate_sequential_number<T>(num_docs_ / 2);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GT);
ASSERT_TRUE(res);
for (uint32_t id = num_docs_ / 2 + 1; id < num_docs_; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
ASSERT_FALSE(res->contains(num_docs_ / 2));
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GE);
ASSERT_TRUE(res);
for (uint32_t id = num_docs_ / 2; id < num_docs_ / 2; ++id) {
if (include_nulls && id % 100 == 0) {
ASSERT_FALSE(res->contains(id));
} else {
ASSERT_TRUE(res->contains(id));
}
}
ASSERT_FALSE(res->contains(num_docs_ / 2 - 1));
v = generate_sequential_number<T>(num_docs_ - 1);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GT);
ASSERT_TRUE(res);
ASSERT_FALSE(res->contains(num_docs_ - 1));
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GE);
ASSERT_TRUE(res);
ASSERT_TRUE(res->contains(num_docs_ - 1));
ASSERT_FALSE(res->contains(num_docs_));
v = generate_sequential_number<T>(num_docs_);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GT);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
res = indexer->search(std::string((char *)&v, sizeof(T)), CompareOp::GE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
}
template <typename T>
void verify_sequential_numbers_null(InvertedColumnIndexer::Ptr indexer) {
InvertedSearchResult::Ptr res = indexer->search_null();
ASSERT_TRUE(res);
for (uint32_t i = 0; i < num_docs_; ++i) {
if (i % 100 == 0) {
ASSERT_TRUE(res->contains(i));
} else {
ASSERT_FALSE(res->contains(i));
}
}
res = indexer->search_non_null();
ASSERT_TRUE(res);
for (uint32_t i = 0; i < num_docs_; ++i) {
if (i % 100 == 0) {
ASSERT_FALSE(res->contains(i));
} else {
ASSERT_TRUE(res->contains(i));
}
}
}
template <typename T>
void verify_sequential_numbers_range_ratio(InvertedColumnIndexer::Ptr indexer,
bool include_nulls) {
uint64_t total_size, range_size;
T v = generate_sequential_number<T>(num_docs_ / 10);
auto s = indexer->evaluate_ratio(std::string((char *)&v, sizeof(T)),
CompareOp::LT, &total_size, &range_size);
ASSERT_TRUE(s.ok());
if (include_nulls) {
ASSERT_EQ(total_size, num_docs_ - num_docs_ / 100);
ASSERT_LE(range_size, num_docs_ / 10 * 2);
} else {
ASSERT_EQ(total_size, num_docs_);
ASSERT_LE(range_size, num_docs_ / 10 * 2);
}
s = indexer->evaluate_ratio(std::string((char *)&v, sizeof(T)),
CompareOp::GT, &total_size, &range_size);
ASSERT_TRUE(s.ok());
if (include_nulls) {
ASSERT_EQ(total_size, num_docs_ - num_docs_ / 100);
ASSERT_GE(range_size, num_docs_ / 10 * 8);
} else {
ASSERT_EQ(total_size, num_docs_);
ASSERT_GE(range_size, num_docs_ / 10 * 8);
}
}
private:
template <typename T>
T generate_sequential_number(uint32_t doc_id) {
// E.g., for int32_t, [id: 5, value: 5]; for float, [id: 5, value: 5.333]
double num_double = doc_id + 0.333;
T num = num_double;
return num;
}
private:
const uint32_t num_docs_;
const uint32_t num_write_threads_;
};
/**
*
* @brief Unit tests for the InvertedColumnIndexer implementation.
*
*/
class InvertedIndexTest : public testing::Test {
/***** Global initialization and cleanup - Start *****/
public:
static void SetUpTestCase() {
zvec::test_util::RemoveTestPath(working_dir);
indexer_ = InvertedIndexer::CreateAndOpen(collection_name, working_dir,
true, {}, false);
params_ = std::make_shared<InvertIndexParams>(true);
}
static void TearDownTestCase() {
indexer_.reset();
zvec::test_util::RemoveTestPath(working_dir);
}
/***** Global initialization and cleanup - End *****/
/***** Per-test initialization and cleanup - Start *****/
protected:
void SetUp() override {}
void TearDown() override {}
/***** Per-test initialization and cleanup - End *****/
protected:
static InvertedIndexer::Ptr indexer_;
static TestHelper test_helper_;
static IndexParams::Ptr params_;
};
InvertedIndexer::Ptr InvertedIndexTest::indexer_{nullptr};
TestHelper InvertedIndexTest::test_helper_{100000, 10};
IndexParams::Ptr InvertedIndexTest::params_{nullptr};
/*
*
* Test Cases
*
*/
TEST_F(InvertedIndexTest, SEQUENTIAL_NUMBERS_INT32) {
ASSERT_TRUE(indexer_);
FieldSchema seq_int32{"seq_int32", DataType::INT32, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(seq_int32).ok());
auto indexer_int32 = (*indexer_)["seq_int32"];
ASSERT_TRUE(indexer_int32);
test_helper_.insert_sequential_numbers<int32_t>(indexer_int32, false);
test_helper_.verify_sequential_numbers<int32_t>(indexer_int32, false);
FieldSchema seq_int32_w_null{"seq_int32_w_null", DataType::INT32, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(seq_int32_w_null).ok());
auto indexer_int32_w_null = (*indexer_)["seq_int32_w_null"];
ASSERT_TRUE(indexer_int32_w_null);
test_helper_.insert_sequential_numbers<int32_t>(indexer_int32_w_null, true);
test_helper_.verify_sequential_numbers<int32_t>(indexer_int32_w_null, true);
}
TEST_F(InvertedIndexTest, SEQUENTIAL_NUMBERS_INT64) {
ASSERT_TRUE(indexer_);
FieldSchema seq_int64{"seq_int64", DataType::INT64, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(seq_int64).ok());
auto indexer_int64 = (*indexer_)["seq_int64"];
ASSERT_TRUE(indexer_int64);
test_helper_.insert_sequential_numbers<int64_t>(indexer_int64, false);
test_helper_.verify_sequential_numbers<int64_t>(indexer_int64, false);
FieldSchema seq_int64_w_null{"seq_int64_w_null", DataType::INT64, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(seq_int64_w_null).ok());
auto indexer_int64_w_null = (*indexer_)["seq_int64_w_null"];
ASSERT_TRUE(indexer_int64_w_null);
test_helper_.insert_sequential_numbers<int64_t>(indexer_int64_w_null, true);
test_helper_.verify_sequential_numbers<int64_t>(indexer_int64_w_null, true);
}
TEST_F(InvertedIndexTest, SEQUENTIAL_NUMBERS_UINT32) {
ASSERT_TRUE(indexer_);
FieldSchema seq_uint32{"seq_uint32", DataType::UINT32, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(seq_uint32).ok());
auto indexer_uint32 = (*indexer_)["seq_uint32"];
ASSERT_TRUE(indexer_uint32);
test_helper_.insert_sequential_numbers<uint32_t>(indexer_uint32, false);
test_helper_.verify_sequential_numbers<uint32_t>(indexer_uint32, false);
FieldSchema seq_uint32_w_null{"seq_uint32_w_null", DataType::UINT32, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(seq_uint32_w_null).ok());
auto indexer_uint32_w_null = (*indexer_)["seq_uint32_w_null"];
ASSERT_TRUE(indexer_uint32_w_null);
test_helper_.insert_sequential_numbers<uint32_t>(indexer_uint32_w_null, true);
test_helper_.verify_sequential_numbers<uint32_t>(indexer_uint32_w_null, true);
}
TEST_F(InvertedIndexTest, SEQUENTIAL_NUMBERS_UINT64) {
ASSERT_TRUE(indexer_);
FieldSchema seq_uint64{"seq_uint64", DataType::UINT64, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(seq_uint64).ok());
auto indexer_uint64 = (*indexer_)["seq_uint64"];
ASSERT_TRUE(indexer_uint64);
test_helper_.insert_sequential_numbers<uint64_t>(indexer_uint64, false);
test_helper_.verify_sequential_numbers<uint64_t>(indexer_uint64, false);
FieldSchema seq_uint64_w_null{"seq_uint64_w_null", DataType::UINT64, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(seq_uint64_w_null).ok());
auto indexer_uint64_w_null = (*indexer_)["seq_uint64_w_null"];
ASSERT_TRUE(indexer_uint64_w_null);
test_helper_.insert_sequential_numbers<uint64_t>(indexer_uint64_w_null, true);
test_helper_.verify_sequential_numbers<uint64_t>(indexer_uint64_w_null, true);
}
TEST_F(InvertedIndexTest, SEQUENTIAL_NUMBERS_FLOAT) {
ASSERT_TRUE(indexer_);
FieldSchema seq_float{"seq_float", DataType::FLOAT, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(seq_float).ok());
auto indexer_float = (*indexer_)["seq_float"];
ASSERT_TRUE(indexer_float);
test_helper_.insert_sequential_numbers<float>(indexer_float, false);
test_helper_.verify_sequential_numbers<float>(indexer_float, false);
FieldSchema seq_float_w_null{"seq_float_w_null", DataType::FLOAT, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(seq_float_w_null).ok());
auto indexer_float_w_null = (*indexer_)["seq_float_w_null"];
ASSERT_TRUE(indexer_float_w_null);
test_helper_.insert_sequential_numbers<float>(indexer_float_w_null, true);
test_helper_.verify_sequential_numbers<float>(indexer_float_w_null, true);
}
TEST_F(InvertedIndexTest, SEQUENTIAL_NUMBERS_DOUBLE) {
ASSERT_TRUE(indexer_);
FieldSchema seq_double{"seq_double", DataType::DOUBLE, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(seq_double).ok());
auto indexer_double = (*indexer_)["seq_double"];
ASSERT_TRUE(indexer_double);
test_helper_.insert_sequential_numbers<double>(indexer_double, false);
test_helper_.verify_sequential_numbers<double>(indexer_double, false);
FieldSchema seq_double_w_null{"seq_double_w_null", DataType::DOUBLE, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(seq_double_w_null).ok());
auto indexer_double_w_null = (*indexer_)["seq_double_w_null"];
ASSERT_TRUE(indexer_double_w_null);
test_helper_.insert_sequential_numbers<double>(indexer_double_w_null, true);
test_helper_.verify_sequential_numbers<double>(indexer_double_w_null, true);
}
TEST_F(InvertedIndexTest, SEALED) {
ASSERT_TRUE(indexer_);
ASSERT_TRUE(indexer_->seal().ok());
auto indexer_int32 = (*indexer_)["seq_int32"];
ASSERT_TRUE(indexer_int32);
test_helper_.verify_sequential_numbers<int32_t>(indexer_int32, false);
auto indexer_int32_w_null = (*indexer_)["seq_int32_w_null"];
ASSERT_TRUE(indexer_int32_w_null);
test_helper_.verify_sequential_numbers<int32_t>(indexer_int32_w_null, true);
auto indexer_int64 = (*indexer_)["seq_int64"];
ASSERT_TRUE(indexer_int64);
test_helper_.verify_sequential_numbers<int64_t>(indexer_int64, false);
auto indexer_int64_w_null = (*indexer_)["seq_int64_w_null"];
ASSERT_TRUE(indexer_int64_w_null);
test_helper_.verify_sequential_numbers<int64_t>(indexer_int64_w_null, true);
auto indexer_uint32 = (*indexer_)["seq_uint32"];
ASSERT_TRUE(indexer_uint32);
test_helper_.verify_sequential_numbers<uint32_t>(indexer_uint32, false);
auto indexer_uint32_w_null = (*indexer_)["seq_uint32_w_null"];
ASSERT_TRUE(indexer_uint32_w_null);
test_helper_.verify_sequential_numbers<uint32_t>(indexer_uint32_w_null, true);
auto indexer_uint64 = (*indexer_)["seq_uint64"];
ASSERT_TRUE(indexer_uint64);
test_helper_.verify_sequential_numbers<uint64_t>(indexer_uint64, false);
auto indexer_uint64_w_null = (*indexer_)["seq_uint64_w_null"];
ASSERT_TRUE(indexer_uint64_w_null);
test_helper_.verify_sequential_numbers<uint64_t>(indexer_uint64_w_null, true);
auto indexer_float = (*indexer_)["seq_float"];
ASSERT_TRUE(indexer_float);
test_helper_.verify_sequential_numbers<float>(indexer_float, false);
auto indexer_float_w_null = (*indexer_)["seq_float_w_null"];
ASSERT_TRUE(indexer_float_w_null);
test_helper_.verify_sequential_numbers<float>(indexer_float_w_null, true);
auto indexer_double = (*indexer_)["seq_double"];
ASSERT_TRUE(indexer_double);
test_helper_.verify_sequential_numbers<double>(indexer_double, false);
auto indexer_double_w_null = (*indexer_)["seq_double_w_null"];
ASSERT_TRUE(indexer_double_w_null);
test_helper_.verify_sequential_numbers<double>(indexer_double_w_null, true);
}
TEST_F(InvertedIndexTest, CREATE_SNAPSHOT) {
#ifdef __ANDROID__
GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink support (needed by RocksDB checkpoint)";
#endif
ASSERT_TRUE(indexer_);
std::string snapshot_dir = working_dir + "snapshot";
ASSERT_TRUE(indexer_->create_snapshot(snapshot_dir).ok());
std::vector<FieldSchema> fields = {
FieldSchema("seq_int32", DataType::INT32, true, params_),
FieldSchema("seq_int32_w_null", DataType::INT32, true, params_),
FieldSchema("seq_int64", DataType::INT64, true, params_),
FieldSchema("seq_int64_w_null", DataType::INT64, true, params_),
FieldSchema("seq_uint32", DataType::UINT32, true, params_),
FieldSchema("seq_uint32_w_null", DataType::UINT32, true, params_),
FieldSchema("seq_uint64", DataType::UINT64, true, params_),
FieldSchema("seq_uint64_w_null", DataType::UINT64, true, params_),
FieldSchema("seq_float", DataType::FLOAT, true, params_),
FieldSchema("seq_float_w_null", DataType::FLOAT, true, params_),
FieldSchema("seq_double", DataType::DOUBLE, true, params_),
FieldSchema("seq_double_w_null", DataType::DOUBLE, true, params_)};
auto snapshot_indexer = InvertedIndexer::CreateAndOpen(
"snapshot", snapshot_dir, false, fields, false);
ASSERT_TRUE(snapshot_indexer);
auto indexer_int32 = (*snapshot_indexer)["seq_int32"];
ASSERT_TRUE(indexer_int32);
test_helper_.verify_sequential_numbers<int32_t>(indexer_int32, false);
auto indexer_int32_w_null = (*snapshot_indexer)["seq_int32_w_null"];
ASSERT_TRUE(indexer_int32_w_null);
test_helper_.verify_sequential_numbers<int32_t>(indexer_int32_w_null, true);
auto indexer_int64 = (*snapshot_indexer)["seq_int64"];
ASSERT_TRUE(indexer_int64);
test_helper_.verify_sequential_numbers<int64_t>(indexer_int64, false);
auto indexer_int64_w_null = (*snapshot_indexer)["seq_int64_w_null"];
ASSERT_TRUE(indexer_int64_w_null);
test_helper_.verify_sequential_numbers<int64_t>(indexer_int64_w_null, true);
auto indexer_uint32 = (*snapshot_indexer)["seq_uint32"];
ASSERT_TRUE(indexer_uint32);
test_helper_.verify_sequential_numbers<uint32_t>(indexer_uint32, false);
auto indexer_uint32_w_null = (*snapshot_indexer)["seq_uint32_w_null"];
ASSERT_TRUE(indexer_uint32_w_null);
test_helper_.verify_sequential_numbers<uint32_t>(indexer_uint32_w_null, true);
auto indexer_uint64 = (*snapshot_indexer)["seq_uint64"];
ASSERT_TRUE(indexer_uint64);
test_helper_.verify_sequential_numbers<uint64_t>(indexer_uint64, false);
auto indexer_uint64_w_null = (*snapshot_indexer)["seq_uint64_w_null"];
ASSERT_TRUE(indexer_uint64_w_null);
test_helper_.verify_sequential_numbers<uint64_t>(indexer_uint64_w_null, true);
auto indexer_float = (*snapshot_indexer)["seq_float"];
ASSERT_TRUE(indexer_float);
test_helper_.verify_sequential_numbers<float>(indexer_float, false);
auto indexer_float_w_null = (*snapshot_indexer)["seq_float_w_null"];
ASSERT_TRUE(indexer_float_w_null);
test_helper_.verify_sequential_numbers<float>(indexer_float_w_null, true);
auto indexer_double = (*snapshot_indexer)["seq_double"];
ASSERT_TRUE(indexer_double);
test_helper_.verify_sequential_numbers<double>(indexer_double, false);
auto indexer_double_w_null = (*snapshot_indexer)["seq_double_w_null"];
ASSERT_TRUE(indexer_double_w_null);
test_helper_.verify_sequential_numbers<double>(indexer_double_w_null, true);
}
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif
@@ -0,0 +1,378 @@
// 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 <random>
#include <gtest/gtest.h>
#include "db/index/column/inverted_column/inverted_indexer.h"
#include "tests/test_util.h"
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
#endif
using namespace zvec;
using File = ailego::File;
const std::string working_dir{"./inverted_column_indexer_string_dir/"};
const std::string collection_name{"test_collection"};
/**
* @brief A helper class for testing the InvertedColumnIndexer implementation.
*
* This class generates test data with specific patterns to verify the
* correctness of the inverted index implementation. It provides various methods
* to populate an InvertedColumnIndexer with predictable data patterns and
* verify that the indexing and search operations work correctly.
*
*/
class TestHelper {
public:
TestHelper(uint32_t num_docs, uint32_t num_write_threads = 10)
: num_docs_(num_docs / 100 * 100),
num_write_threads_(num_write_threads) {};
void insert_strings(InvertedColumnIndexer::Ptr indexer) {
auto insert_func = [&](uint32_t start, uint32_t end) {
Status s;
for (uint32_t i = start; i < end; ++i) {
auto v = generate_string(i);
s = indexer->insert(i, v);
ASSERT_TRUE(s.ok());
}
};
uint32_t num_docs_per_thread = num_docs_ / num_write_threads_;
std::vector<std::thread> threads{};
for (uint32_t t = 0; t < num_write_threads_; ++t) {
threads.emplace_back(insert_func, t * num_docs_per_thread,
(t + 1) * num_docs_per_thread);
}
for (auto &t : threads) {
t.join();
}
}
void verify_strings(InvertedColumnIndexer::Ptr indexer) {
verify_strings_eq_ne(indexer);
verify_strings_like(indexer);
verify_strings_range(indexer);
}
void verify_strings_eq_ne(InvertedColumnIndexer::Ptr indexer) {
InvertedSearchResult::Ptr res;
// Test EQ operator
for (uint32_t i = 0; i < 20; i++) {
auto v = generate_string(i);
res = indexer->search(v, CompareOp::EQ);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 20);
for (uint32_t j = 0; j < num_docs_ / 20; ++j) {
ASSERT_TRUE(res->contains(i + j * 20));
}
}
// Test NE operator with a non-existent value
std::string v = "NotExist";
res = indexer->search(v, CompareOp::NE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_);
// Test NE operator with a random value
static std::random_device rd;
static std::mt19937 gen(rd());
std::uniform_int_distribution<uint32_t> dis(0, 19);
uint32_t random_num = dis(gen);
v = generate_string(random_num);
res = indexer->search(v, CompareOp::NE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ - num_docs_ / 20);
for (uint32_t j = 0; j < num_docs_; ++j) {
if (j % 20 == random_num) {
ASSERT_FALSE(res->contains(j));
} else {
ASSERT_TRUE(res->contains(j));
}
}
}
void verify_strings_like(InvertedColumnIndexer::Ptr indexer) {
InvertedSearchResult::Ptr res;
std::string v = "Three";
res = indexer->search(v, CompareOp::HAS_PREFIX);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 4);
for (uint32_t j = 0; j < num_docs_; ++j) {
if (j % 4 == 2) {
ASSERT_TRUE(res->contains(j));
} else {
ASSERT_FALSE(res->contains(j));
}
}
v = "06";
res = indexer->search(v, CompareOp::HAS_SUFFIX);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 20);
for (uint32_t j = 0; j < num_docs_; ++j) {
if (j % 20 == 6) {
ASSERT_TRUE(res->contains(j));
} else {
ASSERT_FALSE(res->contains(j));
}
}
v = "6";
res = indexer->search(v, CompareOp::HAS_SUFFIX);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ / 10);
for (uint32_t j = 0; j < num_docs_; ++j) {
if (j % 20 == 6 || j % 20 == 16) {
ASSERT_TRUE(res->contains(j));
} else {
ASSERT_FALSE(res->contains(j));
}
}
v = "21";
res = indexer->search(v, CompareOp::HAS_SUFFIX);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
}
void verify_strings_range(InvertedColumnIndexer::Ptr indexer) {
InvertedSearchResult::Ptr res;
std::string v = "Two";
res = indexer->search(v, CompareOp::LT);
ASSERT_TRUE(res);
// "One", "Three", and "Four" are less than "Two" in string sense
ASSERT_EQ(res->count(), num_docs_ / 4 * 3);
for (uint32_t j = 0; j < num_docs_; ++j) {
if (j % 4 == 1) {
ASSERT_FALSE(res->contains(j));
} else {
ASSERT_TRUE(res->contains(j));
}
}
}
void insert_string_arrays(InvertedColumnIndexer::Ptr indexer) {
auto insert_func = [&](uint32_t start, uint32_t end) {
Status s;
for (uint32_t i = start; i < end; ++i) {
auto v = generate_string_array(i);
s = indexer->insert(i, v);
ASSERT_TRUE(s.ok());
}
};
uint32_t num_docs_per_thread = num_docs_ / num_write_threads_;
std::vector<std::thread> threads{};
for (uint32_t t = 0; t < num_write_threads_; ++t) {
threads.emplace_back(insert_func, t * num_docs_per_thread,
(t + 1) * num_docs_per_thread);
}
for (auto &t : threads) {
t.join();
}
}
void verify_string_arrays(InvertedColumnIndexer::Ptr indexer) {
InvertedSearchResult::Ptr res;
auto v = generate_string_array(100);
res = indexer->multi_search(v, CompareOp::CONTAIN_ALL);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 1);
ASSERT_TRUE(res->contains(100));
res = indexer->multi_search(v, CompareOp::CONTAIN_ANY);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 5);
ASSERT_TRUE(res->contains(98));
ASSERT_TRUE(res->contains(99));
ASSERT_TRUE(res->contains(100));
ASSERT_TRUE(res->contains(101));
ASSERT_TRUE(res->contains(102));
res = indexer->multi_search(v, CompareOp::NOT_CONTAIN_ALL);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ - 1);
ASSERT_FALSE(res->contains(100));
res = indexer->multi_search(v, CompareOp::NOT_CONTAIN_ANY);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_ - 5);
ASSERT_FALSE(res->contains(98));
ASSERT_FALSE(res->contains(99));
ASSERT_FALSE(res->contains(100));
ASSERT_FALSE(res->contains(101));
ASSERT_FALSE(res->contains(102));
res = indexer->search_array_len(3, CompareOp::EQ);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), num_docs_);
res = indexer->search_array_len(3, CompareOp::NE);
ASSERT_TRUE(res);
ASSERT_EQ(res->count(), 0);
}
private:
std::string generate_string(uint32_t doc_id) {
std::string prefix;
switch (doc_id % 4) {
case 0:
prefix = "One";
break;
case 1:
prefix = "Two";
break;
case 2:
prefix = "Three";
break;
case 3:
prefix = "Four";
break;
}
std::stringstream suffix;
suffix << std::setfill('0') << std::setw(2) << doc_id % 20;
return prefix + "_" + suffix.str();
}
std::vector<std::string> generate_string_array(uint32_t doc_id) {
std::vector<std::string> ret;
std::stringstream ss1;
ss1 << std::setfill('0') << std::setw(10) << doc_id;
ret.emplace_back(ss1.str());
std::stringstream ss2;
ss2 << std::setfill('0') << std::setw(10) << doc_id + 1;
ret.emplace_back(ss2.str());
std::stringstream ss3;
ss3 << std::setfill('0') << std::setw(10) << doc_id + 2;
ret.emplace_back(ss3.str());
return ret;
}
private:
const uint32_t num_docs_;
const uint32_t num_write_threads_;
};
/**
*
* @brief Unit tests for the InvertedColumnIndexer implementation.
*
*/
class InvertedIndexTest : public testing::Test {
/***** Global initialization and cleanup - Start *****/
public:
static void SetUpTestCase() {
zvec::test_util::RemoveTestPath(working_dir);
indexer_ = InvertedIndexer::CreateAndOpen(collection_name, working_dir,
true, {}, false);
params_ = std::make_shared<InvertIndexParams>(true, true);
}
static void TearDownTestCase() {
indexer_.reset();
zvec::test_util::RemoveTestPath(working_dir);
}
/***** Global initialization and cleanup - End *****/
/***** Per-test initialization and cleanup - Start *****/
protected:
void SetUp() override {}
void TearDown() override {}
/***** Per-test initialization and cleanup - End *****/
protected:
static InvertedIndexer::Ptr indexer_;
static TestHelper test_helper_;
static IndexParams::Ptr params_;
};
InvertedIndexer::Ptr InvertedIndexTest::indexer_{nullptr};
TestHelper InvertedIndexTest::test_helper_{100000, 10};
IndexParams::Ptr InvertedIndexTest::params_{nullptr};
/*
*
* Test Cases
*
*/
TEST_F(InvertedIndexTest, STRINGS) {
ASSERT_TRUE(indexer_);
FieldSchema test_string{"test_string", DataType::STRING, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(test_string).ok());
auto indexer_string = (*indexer_)["test_string"];
ASSERT_TRUE(indexer_string);
test_helper_.insert_strings(indexer_string);
test_helper_.verify_strings(indexer_string);
}
TEST_F(InvertedIndexTest, STRING_ARRAYS) {
ASSERT_TRUE(indexer_);
FieldSchema test_string_array{"test_string_array", DataType::ARRAY_STRING,
true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(test_string_array).ok());
auto indexer_string_array = (*indexer_)["test_string_array"];
ASSERT_TRUE(indexer_string_array);
test_helper_.insert_string_arrays(indexer_string_array);
test_helper_.verify_string_arrays(indexer_string_array);
}
TEST_F(InvertedIndexTest, SEALED) {
ASSERT_TRUE(indexer_);
ASSERT_TRUE(indexer_->seal().ok());
auto indexer_string = (*indexer_)["test_string"];
ASSERT_TRUE(indexer_string);
test_helper_.verify_strings(indexer_string);
auto indexer_string_array = (*indexer_)["test_string_array"];
ASSERT_TRUE(indexer_string_array);
test_helper_.verify_string_arrays(indexer_string_array);
}
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif
@@ -0,0 +1,256 @@
// 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 <gtest/gtest.h>
#include "tests/test_util.h"
#define private public
#define protected public
#include "db/index/column/inverted_column/inverted_indexer.h"
#undef private
#undef protected
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
#endif
using namespace zvec;
using File = ailego::File;
const std::string working_dir{"./inverted_indexer_util_dir/"};
const std::string collection_name{"test_collection"};
class InvertedIndexTest : public testing::Test {
/***** Global initialization and cleanup - Start *****/
public:
static void SetUpTestCase() {
zvec::test_util::RemoveTestPath(working_dir);
indexer_ = InvertedIndexer::CreateAndOpen(collection_name, working_dir,
true, {}, false);
params_ = std::make_shared<InvertIndexParams>(true, false);
}
static void TearDownTestCase() {
indexer_.reset();
zvec::test_util::RemoveTestPath(working_dir);
}
/***** Global initialization and cleanup - End *****/
/***** Per-test initialization and cleanup - Start *****/
protected:
void SetUp() override {}
void TearDown() override {}
/***** Per-test initialization and cleanup - End *****/
protected:
static InvertedIndexer::Ptr indexer_;
static IndexParams::Ptr params_;
};
InvertedIndexer::Ptr InvertedIndexTest::indexer_{nullptr};
IndexParams::Ptr InvertedIndexTest::params_{nullptr};
TEST_F(InvertedIndexTest, COLLECTION_NAME) {
ASSERT_TRUE(indexer_);
ASSERT_EQ(indexer_->collection(), collection_name);
}
TEST_F(InvertedIndexTest, WORKING_DIR) {
ASSERT_TRUE(indexer_);
ASSERT_EQ(indexer_->working_dir(), working_dir);
}
TEST_F(InvertedIndexTest, COLUMN_MANIPULATION_EDGE_CASE) {
ASSERT_FALSE(indexer_->remove_column_indexer("Non-exist").ok());
FieldSchema field{"field_int32", DataType::INT32, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(field).ok());
auto indexer_int32 = (*indexer_)["field_int32"];
ASSERT_NE(indexer_int32, nullptr);
FieldSchema field_duplicate{"field_int32", DataType::INT32, false, params_};
ASSERT_FALSE(indexer_->create_column_indexer(field_duplicate).ok());
ASSERT_TRUE(indexer_->remove_column_indexer("field_int32").ok());
}
TEST_F(InvertedIndexTest, COLUMN_MANIPULATION_INT32) {
ASSERT_TRUE(indexer_);
// Create column indexer
FieldSchema field{"field_int32", DataType::INT32, true, params_};
ASSERT_TRUE(indexer_->create_column_indexer(field).ok());
auto indexer_int32 = (*indexer_)["field_int32"];
ASSERT_NE(indexer_int32, nullptr);
// Insert some data
int32_t i;
for (i = 0; i < 3000; i++) {
auto s = indexer_int32->insert(i, std::string((char *)&i, sizeof(int32_t)));
ASSERT_TRUE(s.ok());
}
// Store variable names for later retrieval
auto cf_name_terms = indexer_int32->cf_name_terms();
auto cf_name_ranges = indexer_int32->cf_name_ranges();
auto cf_name_cdf = indexer_int32->cf_name_cdf();
auto key_max_id = indexer_int32->key_max_id();
auto key_null = indexer_int32->key_null();
auto key_sealed = indexer_int32->key_sealed();
ASSERT_TRUE(indexer_int32->seal().ok());
auto s = indexer_int32->insert(i, std::string((char *)&i, sizeof(int32_t)));
ASSERT_FALSE(s.ok());
std::string value;
ASSERT_TRUE(indexer_->rocksdb_context_.db_->Get({}, key_max_id, &value).ok());
ASSERT_TRUE(indexer_->rocksdb_context_.db_->Get({}, key_sealed, &value).ok());
// Remove column indexer
ASSERT_TRUE(indexer_->remove_column_indexer("field_int32").ok());
indexer_int32 = (*indexer_)["field_int32"];
ASSERT_EQ(indexer_int32, nullptr);
// No garbage left
ASSERT_EQ(indexer_->rocksdb_context_.get_cf(cf_name_terms), nullptr);
ASSERT_EQ(indexer_->rocksdb_context_.get_cf(cf_name_ranges), nullptr);
auto cdf = indexer_->rocksdb_context_.get_cf(cf_name_cdf);
ASSERT_NE(cdf, nullptr);
ASSERT_EQ(
indexer_->rocksdb_context_.db_->Get({}, cdf, field.name(), &value).code(),
rocksdb::Status::kNotFound);
ASSERT_EQ(indexer_->rocksdb_context_.db_->Get({}, key_max_id, &value).code(),
rocksdb::Status::kNotFound);
ASSERT_EQ(indexer_->rocksdb_context_.db_->Get({}, key_null, &value).code(),
rocksdb::Status::kNotFound);
ASSERT_EQ(indexer_->rocksdb_context_.db_->Get({}, key_sealed, &value).code(),
rocksdb::Status::kNotFound);
}
TEST_F(InvertedIndexTest, COLUMN_MANIPULATION_ARRAY_STRING) {
ASSERT_TRUE(indexer_);
// Create column indexer
FieldSchema field{"field_string_array", DataType::ARRAY_STRING, true,
params_};
ASSERT_TRUE(indexer_->create_column_indexer(field).ok());
auto indexer_string_array = (*indexer_)["field_string_array"];
ASSERT_NE(indexer_string_array, nullptr);
// Insert some data
for (uint32_t i = 0; i < 1500; i++) {
std::vector<std::string> values;
for (uint32_t j = 0; j < 5; j++) {
values.emplace_back("Number_" + std::to_string(i));
}
auto s = indexer_string_array->insert(i, values);
ASSERT_TRUE(s.ok());
}
// Store variable names for later retrieval
auto cf_name_terms = indexer_string_array->cf_name_terms();
auto cf_name_array_len = indexer_string_array->cf_name_array_len();
auto cf_name_ranges = indexer_string_array->cf_name_ranges();
ASSERT_EQ(indexer_->rocksdb_context_.get_cf(cf_name_ranges), nullptr);
auto cf_name_cdf = indexer_string_array->cf_name_cdf();
auto key_max_id = indexer_string_array->key_max_id();
auto key_null = indexer_string_array->key_null();
auto key_sealed = indexer_string_array->key_sealed();
// Remove column indexer
ASSERT_TRUE(indexer_->remove_column_indexer("field_string_array").ok());
indexer_string_array = (*indexer_)["field_string_array"];
ASSERT_EQ(indexer_string_array, nullptr);
// No garbage left
std::string value;
ASSERT_EQ(indexer_->rocksdb_context_.get_cf(cf_name_terms), nullptr);
ASSERT_EQ(indexer_->rocksdb_context_.get_cf(cf_name_array_len), nullptr);
ASSERT_EQ(indexer_->rocksdb_context_.get_cf(cf_name_ranges), nullptr);
auto cdf = indexer_->rocksdb_context_.get_cf(cf_name_cdf);
ASSERT_NE(cdf, nullptr);
ASSERT_EQ(
indexer_->rocksdb_context_.db_->Get({}, cdf, field.name(), &value).code(),
rocksdb::Status::kNotFound);
ASSERT_EQ(indexer_->rocksdb_context_.db_->Get({}, key_max_id, &value).code(),
rocksdb::Status::kNotFound);
ASSERT_EQ(indexer_->rocksdb_context_.db_->Get({}, key_null, &value).code(),
rocksdb::Status::kNotFound);
ASSERT_EQ(indexer_->rocksdb_context_.db_->Get({}, key_sealed, &value).code(),
rocksdb::Status::kNotFound);
}
TEST_F(InvertedIndexTest, INVERTED_SEARCH_RESULT) {
roaring_bitmap_t *bitmap1 = roaring_bitmap_create();
roaring_bitmap_add(bitmap1, 1);
roaring_bitmap_add(bitmap1, 2);
roaring_bitmap_add(bitmap1, 3);
auto res1 = std::make_shared<InvertedSearchResult>(bitmap1);
std::vector<uint32_t> ids;
res1->extract_ids(&ids);
ASSERT_EQ(ids.size(), 3);
ASSERT_EQ(ids[0], 1);
ASSERT_EQ(ids[1], 2);
ASSERT_EQ(ids[2], 3);
roaring_bitmap_t *bitmap2 = roaring_bitmap_create();
roaring_bitmap_add(bitmap2, 3);
roaring_bitmap_add(bitmap2, 4);
roaring_bitmap_add(bitmap2, 5);
auto res2 = std::make_shared<InvertedSearchResult>(bitmap2);
res1->AND(*res2);
ASSERT_EQ(res1->count(), 1);
auto filter = res1->make_filter();
ASSERT_TRUE(filter);
ASSERT_FALSE(filter->is_filtered(3));
roaring_bitmap_t *bitmap3 = roaring_bitmap_create();
roaring_bitmap_add(bitmap3, 1);
roaring_bitmap_add(bitmap3, 3);
roaring_bitmap_add(bitmap3, 9);
roaring_bitmap_add(bitmap3, 11);
auto res3 = std::make_shared<InvertedSearchResult>(bitmap3);
res2->OR(*res3);
ASSERT_EQ(res2->count(), 6);
filter = res2->make_filter();
ASSERT_FALSE(filter->is_filtered(1));
ASSERT_FALSE(filter->is_filtered(3));
ASSERT_FALSE(filter->is_filtered(4));
ASSERT_FALSE(filter->is_filtered(5));
ASSERT_FALSE(filter->is_filtered(9));
ASSERT_FALSE(filter->is_filtered(11));
}
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif
@@ -0,0 +1,378 @@
// 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 <cstdint>
#include <numeric>
#include <set>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "db/index/column/vector_column/combined_vector_column_indexer.h"
#include "db/index/column/vector_column/vector_column_indexer.h"
#include "db/index/column/vector_column/vector_column_params.h"
#include "tests/test_util.h"
#include "zvec/db/index_params.h"
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
#endif
using namespace zvec;
namespace {
constexpr uint32_t kGbDimension = 4;
constexpr uint32_t kGbNumDocs = 12;
constexpr uint32_t kGbNumGroups = 3;
constexpr uint32_t kGbGroupTopk = 2;
constexpr uint32_t kGbSearchTopk = 100;
constexpr uint32_t kGbSparseCount = 5;
struct GroupByCase {
std::string name;
IndexParams::Ptr index_params;
QueryParams::Ptr query_params; // core-layer query params; nullptr for sparse
bool is_sparse = false;
uint32_t dimension = kGbDimension;
bool optional = false; // skip when plugin unavailable (e.g. DiskAnn)
bool with_bf_pks = false;
bool fetch_vector = false;
};
std::unique_ptr<vector_column_params::GroupByParams> MakeGroupByParams(
uint32_t group_count = kGbNumGroups) {
return std::make_unique<vector_column_params::GroupByParams>(
kGbGroupTopk, group_count, [](uint64_t key) -> std::string {
return std::to_string(key % kGbNumGroups);
});
}
std::unique_ptr<vector_column_params::GroupByParams> MakeSegmentGroupByParams(
uint32_t group_topk, uint32_t group_count) {
return std::make_unique<vector_column_params::GroupByParams>(
group_topk, group_count,
[](uint64_t key) -> std::string { return key < 2 ? "low" : "high"; });
}
std::vector<uint64_t> AllPks() {
std::vector<uint64_t> pks(kGbNumDocs);
std::iota(pks.begin(), pks.end(), 0ull);
return pks;
}
class GroupByIndexerTest : public ::testing::Test {
protected:
void RunOk(const GroupByCase &tc) {
Run(tc, /*expect_error=*/false);
}
void RunRejected(const GroupByCase &tc) {
Run(tc, /*expect_error=*/true);
}
private:
struct QueryHolder {
std::vector<float> dense;
std::vector<uint32_t> sparse_indices;
std::vector<float> sparse_values;
vector_column_params::VectorData data;
};
void Run(const GroupByCase &tc, bool expect_error) {
const std::string path = "test_groupby_" + tc.name + ".index";
zvec::test_util::RemoveTestFiles(path);
auto indexer = OpenIndexer(tc, path);
if (indexer == nullptr) {
zvec::test_util::RemoveTestFiles(path);
return; // optional plugin unavailable
}
InsertDocs(indexer, tc);
QueryHolder holder = MakeQuery(tc);
vector_column_params::QueryParams qp = MakeQueryParams(tc);
auto results = indexer->Search(holder.data, qp);
if (expect_error) {
ASSERT_FALSE(results.has_value())
<< "group_by should be rejected for " << tc.name;
} else {
ASSERT_TRUE(results.has_value()) << tc.name;
AssertGroupedResult(results.value().get(), tc);
}
indexer->Close();
zvec::test_util::RemoveTestFiles(path);
}
static FieldSchema MakeSchema(const GroupByCase &tc) {
if (tc.is_sparse) {
return FieldSchema("test", DataType::SPARSE_VECTOR_FP32, false,
tc.index_params);
}
return FieldSchema("test", DataType::VECTOR_FP32, tc.dimension, false,
tc.index_params);
}
static VectorColumnIndexer::Ptr OpenIndexer(const GroupByCase &tc,
const std::string &path) {
auto indexer = std::make_shared<VectorColumnIndexer>(path, MakeSchema(tc));
if (!indexer->Open(vector_column_params::ReadOptions{true, true}).ok()) {
return nullptr;
}
return indexer;
}
static void InsertDocs(const VectorColumnIndexer::Ptr &indexer,
const GroupByCase &tc) {
for (uint32_t i = 0; i < kGbNumDocs; ++i) {
if (tc.is_sparse) {
std::vector<uint32_t> indices(kGbSparseCount);
std::vector<float> values(kGbSparseCount);
for (uint32_t j = 0; j < kGbSparseCount; ++j) {
indices[j] = i * kGbSparseCount + j;
values[j] = static_cast<float>(i + 1);
}
vector_column_params::SparseVector sv{kGbSparseCount, indices.data(),
values.data()};
ASSERT_TRUE(
indexer->Insert(vector_column_params::VectorData{sv}, i).ok());
} else {
std::vector<float> vec(tc.dimension, static_cast<float>(i));
vector_column_params::DenseVector dv{vec.data()};
ASSERT_TRUE(
indexer->Insert(vector_column_params::VectorData{dv}, i).ok());
}
}
}
static QueryHolder MakeQuery(const GroupByCase &tc) {
QueryHolder h;
if (tc.is_sparse) {
h.sparse_indices.resize(kGbSparseCount);
h.sparse_values.assign(kGbSparseCount, 1.0f);
std::iota(h.sparse_indices.begin(), h.sparse_indices.end(), 0u);
h.data =
vector_column_params::VectorData{vector_column_params::SparseVector{
kGbSparseCount, h.sparse_indices.data(), h.sparse_values.data()}};
} else {
h.dense.assign(tc.dimension, 1.0f);
h.data = vector_column_params::VectorData{
vector_column_params::DenseVector{h.dense.data()}};
}
return h;
}
static vector_column_params::QueryParams MakeQueryParams(
const GroupByCase &tc) {
vector_column_params::QueryParams qp;
qp.topk = kGbSearchTopk;
qp.filter = nullptr;
qp.fetch_vector = tc.fetch_vector;
qp.query_params = tc.query_params;
if (tc.with_bf_pks) {
qp.bf_pks = {AllPks()};
}
qp.group_by = MakeGroupByParams();
return qp;
}
static void AssertGroupedResult(IndexResults *results,
const GroupByCase &tc) {
auto *group_results = dynamic_cast<GroupVectorIndexResults *>(results);
ASSERT_TRUE(group_results)
<< "Expected GroupVectorIndexResults for " << tc.name;
ASSERT_EQ(kGbNumGroups, group_results->groups().size()) << tc.name;
std::set<std::string> group_ids;
for (const auto &group : group_results->groups()) {
group_ids.insert(group.group_id());
ASSERT_LE(group.docs().size(), kGbGroupTopk) << tc.name;
ASSERT_GE(group.docs().size(), 1u) << tc.name;
const uint32_t expected_mod = std::stoul(group.group_id());
for (const auto &doc : group.docs()) {
ASSERT_EQ(expected_mod, doc.key() % kGbNumGroups)
<< tc.name << " doc " << doc.key();
}
for (size_t j = 1; j < group.docs().size(); ++j) {
ASSERT_GE(group.docs()[j - 1].score(), group.docs()[j].score())
<< tc.name << " group " << group.group_id();
}
}
for (uint32_t g = 0; g < kGbNumGroups; ++g) {
ASSERT_TRUE(group_ids.count(std::to_string(g)) > 0)
<< tc.name << " missing group " << g;
}
auto iter = group_results->create_iterator();
size_t total = 0;
while (iter->valid()) {
if (tc.fetch_vector && !tc.is_sparse) {
const auto vector_data = iter->vector();
const auto &dense_vector =
std::get<vector_column_params::DenseVector>(vector_data.vector);
const float *vector =
reinterpret_cast<const float *>(dense_vector.data);
const float expected = static_cast<float>(iter->doc_id());
for (uint32_t i = 0; i < tc.dimension; ++i) {
ASSERT_FLOAT_EQ(expected, vector[i])
<< tc.name << " doc " << iter->doc_id() << " i " << i;
}
}
total++;
iter->next();
}
ASSERT_EQ(group_results->count(), total) << tc.name;
}
};
} // namespace
TEST_F(GroupByIndexerTest, Dense) {
auto hnsw_linear_qp = std::make_shared<HnswQueryParams>(300);
hnsw_linear_qp->set_is_linear(true);
std::vector<GroupByCase> cases{
{"dense_flat_graph", std::make_shared<FlatIndexParams>(MetricType::IP),
std::make_shared<QueryParams>(IndexType::FLAT)},
{"dense_hnsw_graph",
std::make_shared<HnswIndexParams>(MetricType::IP, 10, 100),
std::make_shared<HnswQueryParams>(300)},
{"dense_hnsw_linear",
std::make_shared<HnswIndexParams>(MetricType::IP, 10, 100),
hnsw_linear_qp},
{"dense_hnsw_bf_pks",
std::make_shared<HnswIndexParams>(MetricType::IP, 10, 100),
std::make_shared<HnswQueryParams>(300),
/*is_sparse=*/false, /*dimension=*/kGbDimension,
/*optional=*/false, /*with_bf_pks=*/true},
{"dense_hnsw_fetch_vector",
std::make_shared<HnswIndexParams>(MetricType::IP, 10, 100),
std::make_shared<HnswQueryParams>(300),
/*is_sparse=*/false, /*dimension=*/kGbDimension,
/*optional=*/false, /*with_bf_pks=*/false, /*fetch_vector=*/true},
{"dense_hnsw_fp16_fetch_vector",
std::make_shared<HnswIndexParams>(MetricType::IP, 10, 100,
QuantizeType::FP16),
std::make_shared<HnswQueryParams>(300),
/*is_sparse=*/false, /*dimension=*/kGbDimension,
/*optional=*/false, /*with_bf_pks=*/false, /*fetch_vector=*/true},
};
for (const auto &tc : cases) {
RunOk(tc);
}
}
TEST_F(GroupByIndexerTest, CombinedSortsGroupsBeforeTruncating) {
// Build two blocks where the best group is in the later block; group_count
// must be applied after cross-block group sorting, not merge order.
const std::string block0_path = "test_groupby_combined_block0.index";
const std::string block1_path = "test_groupby_combined_block1.index";
zvec::test_util::RemoveTestFiles(block0_path);
zvec::test_util::RemoveTestFiles(block1_path);
auto index_params = std::make_shared<FlatIndexParams>(MetricType::IP);
FieldSchema schema("test", DataType::VECTOR_FP32, kGbDimension, false,
index_params);
auto block0 = std::make_shared<VectorColumnIndexer>(block0_path, schema);
auto block1 = std::make_shared<VectorColumnIndexer>(block1_path, schema);
ASSERT_TRUE(block0->Open(vector_column_params::ReadOptions{true, true}).ok());
ASSERT_TRUE(block1->Open(vector_column_params::ReadOptions{true, true}).ok());
auto insert_dense = [](const VectorColumnIndexer::Ptr &indexer,
uint32_t doc_id, float value) {
std::vector<float> vec(kGbDimension, value);
vector_column_params::DenseVector dense{vec.data()};
ASSERT_TRUE(
indexer->Insert(vector_column_params::VectorData{dense}, doc_id).ok());
};
insert_dense(block0, 0, 0.0f);
insert_dense(block0, 1, 1.0f);
insert_dense(block1, 0, 10.0f);
insert_dense(block1, 1, 11.0f);
std::vector<BlockMeta> blocks{
BlockMeta(0, BlockType::VECTOR_INDEX, 0, 1, 2, {"test"}),
BlockMeta(1, BlockType::VECTOR_INDEX, 2, 3, 2, {"test"}),
};
SegmentMeta segment_meta;
CombinedVectorColumnIndexer combined({block0, block1}, {}, schema,
segment_meta, blocks, MetricType::IP);
std::vector<float> query(kGbDimension, 1.0f);
vector_column_params::DenseVector dense_query{query.data()};
vector_column_params::QueryParams query_params;
query_params.topk = kGbSearchTopk;
query_params.query_params = std::make_shared<QueryParams>(IndexType::FLAT);
query_params.group_by = MakeSegmentGroupByParams(/*group_topk=*/1,
/*group_count=*/1);
auto results = combined.Search(vector_column_params::VectorData{dense_query},
query_params);
ASSERT_TRUE(results.has_value());
auto *group_results =
dynamic_cast<GroupVectorIndexResults *>(results.value().get());
ASSERT_TRUE(group_results);
ASSERT_EQ(1u, group_results->groups().size());
ASSERT_EQ("high", group_results->groups()[0].group_id());
ASSERT_EQ(1u, group_results->groups()[0].docs().size());
ASSERT_EQ(3u, group_results->groups()[0].docs()[0].key());
ASSERT_FLOAT_EQ(44.0f, group_results->groups()[0].docs()[0].score());
ASSERT_TRUE(block0->Close().ok());
ASSERT_TRUE(block1->Close().ok());
zvec::test_util::RemoveTestFiles(block0_path);
zvec::test_util::RemoveTestFiles(block1_path);
}
TEST_F(GroupByIndexerTest, Sparse) {
std::vector<GroupByCase> cases{
{"sparse_flat_graph", std::make_shared<FlatIndexParams>(MetricType::IP),
/*query_params=*/nullptr,
/*is_sparse=*/true},
{"sparse_hnsw_graph",
std::make_shared<HnswIndexParams>(MetricType::IP, 10, 100),
/*query_params=*/nullptr,
/*is_sparse=*/true},
};
for (const auto &tc : cases) {
RunOk(tc);
}
}
TEST_F(GroupByIndexerTest, UnsupportedIndexTypes) {
std::vector<GroupByCase> cases{
{"unsupported_ivf", std::make_shared<IVFIndexParams>(MetricType::IP, 4),
std::make_shared<IVFQueryParams>(4)},
{"unsupported_diskann",
std::make_shared<DiskAnnIndexParams>(MetricType::IP),
std::make_shared<DiskAnnQueryParams>(),
/*is_sparse=*/false, /*dimension=*/kGbDimension,
/*optional=*/true},
};
for (const auto &tc : cases) {
RunRejected(tc);
}
}
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,550 @@
// 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 <gtest/gtest.h>
#include "db/index/common/proto_converter.h"
#include "db/index/common/type_helper.h"
using namespace zvec;
TEST(ConverterTest, InvertIndexParamsConversion) {
// Test conversion from protobuf to C++ InvertIndexParams
proto::InvertIndexParams invert_pb;
invert_pb.set_enable_range_optimization(true);
auto invert_params = ProtoConverter::FromPb(invert_pb);
ASSERT_NE(invert_params, nullptr);
EXPECT_TRUE(invert_params->enable_range_optimization());
EXPECT_EQ(invert_params->type(), IndexType::INVERT);
// Test with false value
proto::InvertIndexParams invert_pb2;
invert_pb2.set_enable_range_optimization(false);
auto invert_params2 = ProtoConverter::FromPb(invert_pb2);
ASSERT_NE(invert_params2, nullptr);
EXPECT_FALSE(invert_params2->enable_range_optimization());
// Test conversion from C++ to protobuf
InvertIndexParams original_params(true);
auto pb_result = ProtoConverter::ToPb(&original_params);
EXPECT_TRUE(pb_result.enable_range_optimization());
}
TEST(ConverterTest, HnswIndexParamsConversion) {
// Test conversion from protobuf to C++ HnswIndexParams
proto::HnswIndexParams hnsw_pb;
auto *base_params = hnsw_pb.mutable_base();
base_params->set_metric_type(proto::MT_L2);
base_params->set_quantize_type(proto::QT_FP16);
hnsw_pb.set_m(16);
hnsw_pb.set_ef_construction(100);
auto hnsw_params = ProtoConverter::FromPb(hnsw_pb);
ASSERT_NE(hnsw_params, nullptr);
EXPECT_EQ(hnsw_params->metric_type(), MetricType::L2);
EXPECT_EQ(hnsw_params->m(), 16);
EXPECT_EQ(hnsw_params->ef_construction(), 100);
EXPECT_EQ(hnsw_params->quantize_type(), QuantizeType::FP16);
EXPECT_EQ(hnsw_params->type(), IndexType::HNSW);
// Test conversion from C++ to protobuf
HnswIndexParams original_params(MetricType::IP, 32, 200, QuantizeType::INT8);
auto pb_result = ProtoConverter::ToPb(&original_params);
EXPECT_EQ(pb_result.base().metric_type(), proto::MT_IP);
EXPECT_EQ(pb_result.m(), 32);
EXPECT_EQ(pb_result.ef_construction(), 200);
EXPECT_EQ(pb_result.base().quantize_type(), proto::QT_INT8);
}
TEST(ConverterTest, FlatIndexParamsConversion) {
// Test conversion from protobuf to C++ FlatIndexParams
proto::FlatIndexParams flat_pb;
auto *base_params = flat_pb.mutable_base();
base_params->set_metric_type(proto::MT_COSINE);
base_params->set_quantize_type(proto::QT_INT4);
auto flat_params = ProtoConverter::FromPb(flat_pb);
ASSERT_NE(flat_params, nullptr);
EXPECT_EQ(flat_params->metric_type(), MetricType::COSINE);
EXPECT_EQ(flat_params->quantize_type(), QuantizeType::INT4);
EXPECT_EQ(flat_params->type(), IndexType::FLAT);
// Test conversion from C++ to protobuf
FlatIndexParams original_params(MetricType::L2, QuantizeType::FP16);
auto pb_result = ProtoConverter::ToPb(&original_params);
EXPECT_EQ(pb_result.base().metric_type(), proto::MT_L2);
EXPECT_EQ(pb_result.base().quantize_type(), proto::QT_FP16);
}
TEST(ConverterTest, IVFIndexParamsConversion) {
// Test conversion from protobuf to C++ IVFIndexParams
proto::IVFIndexParams ivf_pb;
auto *base_params = ivf_pb.mutable_base();
base_params->set_metric_type(proto::MT_IP);
base_params->set_quantize_type(proto::QT_INT8);
ivf_pb.set_n_list(128);
auto ivf_params = ProtoConverter::FromPb(ivf_pb);
ASSERT_NE(ivf_params, nullptr);
EXPECT_EQ(ivf_params->metric_type(), MetricType::IP);
EXPECT_EQ(ivf_params->n_list(), 128);
EXPECT_EQ(ivf_params->quantize_type(), QuantizeType::INT8);
EXPECT_EQ(ivf_params->type(), IndexType::IVF);
// Test conversion from C++ to protobuf
IVFIndexParams original_params(MetricType::COSINE, 256, 10, false,
QuantizeType::INT4);
auto pb_result = ProtoConverter::ToPb(&original_params);
EXPECT_EQ(pb_result.base().metric_type(), proto::MT_COSINE);
EXPECT_EQ(pb_result.n_list(), 256);
EXPECT_EQ(pb_result.n_iters(), 10);
EXPECT_FALSE(pb_result.use_soar());
EXPECT_EQ(pb_result.base().quantize_type(), proto::QT_INT4);
}
TEST(ConverterTest, IndexParamsConversion) {
// Test conversion from protobuf to C++ IndexParams for HNSW
proto::IndexParams index_pb;
auto *hnsw_pb = index_pb.mutable_hnsw();
auto *base_params = hnsw_pb->mutable_base();
base_params->set_metric_type(proto::MT_L2);
base_params->set_quantize_type(proto::QT_FP16);
hnsw_pb->set_m(16);
hnsw_pb->set_ef_construction(100);
auto index_params = ProtoConverter::FromPb(index_pb);
ASSERT_NE(index_params, nullptr);
EXPECT_EQ(index_params->type(), IndexType::HNSW);
auto hnsw_cast = std::dynamic_pointer_cast<HnswIndexParams>(index_params);
ASSERT_NE(hnsw_cast, nullptr);
EXPECT_EQ(hnsw_cast->metric_type(), MetricType::L2);
EXPECT_EQ(hnsw_cast->m(), 16);
EXPECT_EQ(hnsw_cast->ef_construction(), 100);
EXPECT_EQ(hnsw_cast->quantize_type(), QuantizeType::FP16);
// Test conversion from C++ HnswIndexParams to protobuf IndexParams
HnswIndexParams hnsw_original(MetricType::IP, 32, 200);
auto pb_result = ProtoConverter::ToPb(&hnsw_original);
EXPECT_EQ(pb_result.base().metric_type(), proto::MT_IP);
EXPECT_EQ(pb_result.m(), 32);
EXPECT_EQ(pb_result.ef_construction(), 200);
// Test conversion from protobuf to C++ IndexParams for FLAT
proto::IndexParams index_pb2;
auto *flat_pb = index_pb2.mutable_flat();
auto *base_params2 = flat_pb->mutable_base();
base_params2->set_metric_type(proto::MT_COSINE);
base_params2->set_quantize_type(proto::QT_INT8);
auto index_params2 = ProtoConverter::FromPb(index_pb2);
ASSERT_NE(index_params2, nullptr);
EXPECT_EQ(index_params2->type(), IndexType::FLAT);
auto flat_cast = std::dynamic_pointer_cast<FlatIndexParams>(index_params2);
ASSERT_NE(flat_cast, nullptr);
EXPECT_EQ(flat_cast->metric_type(), MetricType::COSINE);
EXPECT_EQ(flat_cast->quantize_type(), QuantizeType::INT8);
// Test conversion from C++ FlatIndexParams to protobuf IndexParams
FlatIndexParams flat_original(MetricType::L2);
auto pb_result2 = ProtoConverter::ToPb(&flat_original);
EXPECT_EQ(pb_result2.base().metric_type(), proto::MT_L2);
// Test conversion from protobuf to C++ IndexParams for IVF
proto::IndexParams index_pb3;
auto *ivf_pb = index_pb3.mutable_ivf();
auto *base_params3 = ivf_pb->mutable_base();
base_params3->set_metric_type(proto::MT_IP);
base_params3->set_quantize_type(proto::QT_INT4);
ivf_pb->set_n_list(128);
auto index_params3 = ProtoConverter::FromPb(index_pb3);
ASSERT_NE(index_params3, nullptr);
EXPECT_EQ(index_params3->type(), IndexType::IVF);
auto ivf_cast = std::dynamic_pointer_cast<IVFIndexParams>(index_params3);
ASSERT_NE(ivf_cast, nullptr);
EXPECT_EQ(ivf_cast->metric_type(), MetricType::IP);
EXPECT_EQ(ivf_cast->n_list(), 128);
EXPECT_EQ(ivf_cast->quantize_type(), QuantizeType::INT4);
// Test conversion from C++ IVFIndexParams to protobuf IndexParams
IVFIndexParams ivf_original(MetricType::COSINE, 256);
auto pb_result3 = ProtoConverter::ToPb(&ivf_original);
EXPECT_EQ(pb_result3.base().metric_type(), proto::MT_COSINE);
EXPECT_EQ(pb_result3.n_list(), 256);
// Test conversion from protobuf to C++ IndexParams for INVERT
proto::IndexParams index_pb4;
auto *invert_pb = index_pb4.mutable_invert();
invert_pb->set_enable_range_optimization(true);
auto index_params4 = ProtoConverter::FromPb(index_pb4);
ASSERT_NE(index_params4, nullptr);
EXPECT_EQ(index_params4->type(), IndexType::INVERT);
auto invert_cast =
std::dynamic_pointer_cast<InvertIndexParams>(index_params4);
ASSERT_NE(invert_cast, nullptr);
EXPECT_TRUE(invert_cast->enable_range_optimization());
// Test conversion from C++ InvertIndexParams to protobuf IndexParams
InvertIndexParams invert_original(false);
auto pb_result4 = ProtoConverter::ToPb(&invert_original);
EXPECT_FALSE(pb_result4.enable_range_optimization());
}
TEST(ConverterTest, FieldSchemaConversion) {
// Test conversion from protobuf to C++ FieldSchema
proto::FieldSchema field_pb;
field_pb.set_name("test_field");
field_pb.set_data_type(proto::DT_VECTOR_FP32);
field_pb.set_dimension(128);
field_pb.set_nullable(true);
// Add index params
auto *index_params_pb = field_pb.mutable_index_params();
auto *hnsw_pb = index_params_pb->mutable_hnsw();
auto *base_params = hnsw_pb->mutable_base();
base_params->set_metric_type(proto::MT_L2);
base_params->set_quantize_type(proto::QT_FP16);
hnsw_pb->set_m(16);
hnsw_pb->set_ef_construction(100);
auto field_schema = ProtoConverter::FromPb(field_pb);
ASSERT_NE(field_schema, nullptr);
EXPECT_EQ(field_schema->name(), "test_field");
EXPECT_EQ(field_schema->data_type(), DataType::VECTOR_FP32);
EXPECT_TRUE(field_schema->nullable());
EXPECT_EQ(field_schema->dimension(), 128u);
ASSERT_NE(field_schema->index_params(), nullptr);
EXPECT_EQ(field_schema->index_params()->type(), IndexType::HNSW);
// Test conversion from C++ to protobuf
FieldSchema original_field("another_field", DataType::ARRAY_INT32, 64, false,
nullptr);
auto pb_result = ProtoConverter::ToPb(original_field);
EXPECT_EQ(pb_result.name(), "another_field");
EXPECT_EQ(pb_result.data_type(), proto::DT_ARRAY_INT32);
EXPECT_FALSE(pb_result.nullable());
EXPECT_EQ(pb_result.dimension(), 64u);
}
TEST(ConverterTest, CollectionSchemaConversion) {
// Test conversion from protobuf to C++ CollectionSchema
proto::CollectionSchema schema_pb;
schema_pb.set_name("test_collection");
schema_pb.set_max_doc_count_per_segment(1000000);
auto *field1_pb = schema_pb.add_fields();
field1_pb->set_name("field1");
field1_pb->set_data_type(proto::DT_STRING);
auto *field2_pb = schema_pb.add_fields();
field2_pb->set_name("field2");
field2_pb->set_data_type(proto::DT_VECTOR_FP32);
field2_pb->set_dimension(128);
auto collection_schema = ProtoConverter::FromPb(schema_pb);
ASSERT_NE(collection_schema, nullptr);
EXPECT_EQ(collection_schema->name(), "test_collection");
EXPECT_EQ(collection_schema->fields().size(), 2);
EXPECT_EQ(collection_schema->max_doc_count_per_segment(), 1000000u);
// Test conversion from C++ to protobuf
CollectionSchema original_schema;
original_schema.set_name("original_collection");
auto pb_result = ProtoConverter::ToPb(original_schema);
EXPECT_EQ(pb_result.name(), "original_collection");
}
TEST(ConverterTest, BlockMetaConversion) {
// Test conversion from protobuf to C++ BlockMeta
proto::BlockMeta meta_pb;
meta_pb.set_block_id(1);
meta_pb.set_block_type(proto::BT_SCALAR);
meta_pb.set_min_doc_id(100);
meta_pb.set_max_doc_id(200);
meta_pb.set_doc_count(50);
meta_pb.add_columns("col1");
meta_pb.add_columns("col2");
auto block_meta = ProtoConverter::FromPb(meta_pb);
ASSERT_NE(block_meta, nullptr);
EXPECT_EQ(block_meta->id(), 1u);
EXPECT_EQ(block_meta->type(), BlockType::SCALAR);
EXPECT_EQ(block_meta->min_doc_id(), 100u);
EXPECT_EQ(block_meta->max_doc_id(), 200u);
EXPECT_EQ(block_meta->doc_count(), 50u);
EXPECT_EQ(block_meta->columns().size(), 2);
EXPECT_EQ(block_meta->columns()[0], "col1");
EXPECT_EQ(block_meta->columns()[1], "col2");
// Test conversion from C++ to protobuf
BlockMeta original_meta(2, BlockType::VECTOR_INDEX, 300, 400);
original_meta.set_doc_count(75);
original_meta.add_column("col3");
original_meta.add_column("col4");
auto pb_result = ProtoConverter::ToPb(original_meta);
EXPECT_EQ(pb_result.block_id(), 2u);
EXPECT_EQ(pb_result.block_type(), proto::BT_VECTOR_INDEX);
EXPECT_EQ(pb_result.min_doc_id(), 300u);
EXPECT_EQ(pb_result.max_doc_id(), 400u);
EXPECT_EQ(pb_result.doc_count(), 75u);
EXPECT_EQ(pb_result.columns_size(), 2);
EXPECT_EQ(pb_result.columns(0), "col3");
EXPECT_EQ(pb_result.columns(1), "col4");
}
TEST(ConverterTest, SegmentMetaConversion) {
// Test conversion from protobuf to C++ SegmentMeta
proto::SegmentMeta segment_pb;
segment_pb.set_segment_id(10);
// Add persisted blocks
auto *block1_pb = segment_pb.add_persisted_blocks();
block1_pb->set_block_id(1);
block1_pb->set_block_type(proto::BT_SCALAR);
block1_pb->set_min_doc_id(0);
block1_pb->set_max_doc_id(100);
block1_pb->set_doc_count(50);
block1_pb->add_columns("col1");
block1_pb->add_columns("col2");
auto *block2_pb = segment_pb.add_persisted_blocks();
block2_pb->set_block_id(2);
block2_pb->set_block_type(proto::BT_VECTOR_INDEX);
block2_pb->set_min_doc_id(101);
block2_pb->set_max_doc_id(200);
block2_pb->set_doc_count(75);
block2_pb->add_columns("vec_col");
// Add writing forward block
auto *writing_block_pb = segment_pb.mutable_writing_forward_block();
writing_block_pb->set_block_id(3);
writing_block_pb->set_block_type(proto::BT_SCALAR);
writing_block_pb->set_min_doc_id(201);
writing_block_pb->set_max_doc_id(300);
writing_block_pb->set_doc_count(25);
writing_block_pb->add_columns("col3");
// Add indexed vector fields
segment_pb.add_indexed_vector_fields("vec_col1");
segment_pb.add_indexed_vector_fields("vec_col2");
auto segment_meta = ProtoConverter::FromPb(segment_pb);
ASSERT_NE(segment_meta, nullptr);
EXPECT_EQ(segment_meta->id(), 10u);
EXPECT_EQ(segment_meta->persisted_blocks().size(), 2);
EXPECT_TRUE(segment_meta->has_writing_forward_block());
// Check first persisted block
const auto &block1 = segment_meta->persisted_blocks()[0];
EXPECT_EQ(block1.id(), 1u);
EXPECT_EQ(block1.type(), BlockType::SCALAR);
EXPECT_EQ(block1.min_doc_id(), 0u);
EXPECT_EQ(block1.max_doc_id(), 100u);
EXPECT_EQ(block1.doc_count(), 50u);
EXPECT_EQ(block1.columns().size(), 2);
EXPECT_EQ(block1.columns()[0], "col1");
EXPECT_EQ(block1.columns()[1], "col2");
// Check second persisted block
const auto &block2 = segment_meta->persisted_blocks()[1];
EXPECT_EQ(block2.id(), 2u);
EXPECT_EQ(block2.type(), BlockType::VECTOR_INDEX);
EXPECT_EQ(block2.min_doc_id(), 101u);
EXPECT_EQ(block2.max_doc_id(), 200u);
EXPECT_EQ(block2.doc_count(), 75u);
EXPECT_EQ(block2.columns().size(), 1);
EXPECT_EQ(block2.columns()[0], "vec_col");
// Check writing forward block
const auto &writing_block = segment_meta->writing_forward_block();
EXPECT_EQ(writing_block.value().id(), 3u);
EXPECT_EQ(writing_block.value().type(), BlockType::SCALAR);
EXPECT_EQ(writing_block.value().min_doc_id(), 201u);
EXPECT_EQ(writing_block.value().max_doc_id(), 300u);
EXPECT_EQ(writing_block.value().doc_count(), 25u);
EXPECT_EQ(writing_block.value().columns().size(), 1);
EXPECT_EQ(writing_block.value().columns()[0], "col3");
// Check indexed vector fields
EXPECT_TRUE(segment_meta->vector_indexed("vec_col1"));
EXPECT_TRUE(segment_meta->vector_indexed("vec_col2"));
EXPECT_FALSE(segment_meta->vector_indexed("non_existent_field"));
// Test conversion from C++ to protobuf
SegmentMeta original_meta(20);
// Add persisted blocks
BlockMeta block1_meta(1, BlockType::SCALAR_INDEX, 0, 50);
block1_meta.set_doc_count(25);
block1_meta.add_column("col3");
block1_meta.add_column("col4");
original_meta.add_persisted_block(block1_meta);
BlockMeta block2_meta(2, BlockType::VECTOR_INDEX_QUANTIZE, 51, 100);
block2_meta.set_doc_count(30);
block2_meta.add_column("vec_col2");
original_meta.add_persisted_block(block2_meta);
// Set writing forward block
BlockMeta writing_block_meta(3, BlockType::SCALAR, 101, 150);
writing_block_meta.set_doc_count(40);
writing_block_meta.add_column("col5");
original_meta.set_writing_forward_block(writing_block_meta);
// Add indexed vector fields
original_meta.add_indexed_vector_field("vec_field1");
original_meta.add_indexed_vector_field("vec_field2");
auto pb_result = ProtoConverter::ToPb(original_meta);
EXPECT_EQ(pb_result.segment_id(), 20u);
EXPECT_EQ(pb_result.persisted_blocks_size(), 2);
// Check first persisted block
const auto &pb_block1 = pb_result.persisted_blocks(0);
EXPECT_EQ(pb_block1.block_id(), 1u);
EXPECT_EQ(pb_block1.block_type(), proto::BT_SCALAR_INDEX);
EXPECT_EQ(pb_block1.min_doc_id(), 0u);
EXPECT_EQ(pb_block1.max_doc_id(), 50u);
EXPECT_EQ(pb_block1.doc_count(), 25u);
EXPECT_EQ(pb_block1.columns_size(), 2);
EXPECT_EQ(pb_block1.columns(0), "col3");
EXPECT_EQ(pb_block1.columns(1), "col4");
// Check second persisted block
const auto &pb_block2 = pb_result.persisted_blocks(1);
EXPECT_EQ(pb_block2.block_id(), 2u);
EXPECT_EQ(pb_block2.block_type(), proto::BT_VECTOR_INDEX_QUANTIZE);
EXPECT_EQ(pb_block2.min_doc_id(), 51u);
EXPECT_EQ(pb_block2.max_doc_id(), 100u);
EXPECT_EQ(pb_block2.doc_count(), 30u);
EXPECT_EQ(pb_block2.columns_size(), 1);
EXPECT_EQ(pb_block2.columns(0), "vec_col2");
// Check writing forward block
const auto &pb_writing_block = pb_result.writing_forward_block();
EXPECT_EQ(pb_writing_block.block_id(), 3u);
EXPECT_EQ(pb_writing_block.block_type(), proto::BT_SCALAR);
EXPECT_EQ(pb_writing_block.min_doc_id(), 101u);
EXPECT_EQ(pb_writing_block.max_doc_id(), 150u);
EXPECT_EQ(pb_writing_block.doc_count(), 40u);
EXPECT_EQ(pb_writing_block.columns_size(), 1);
EXPECT_EQ(pb_writing_block.columns(0), "col5");
// Check indexed vector fields
EXPECT_EQ(pb_result.indexed_vector_fields_size(), 2);
EXPECT_EQ(pb_result.indexed_vector_fields(0), "vec_field1");
EXPECT_EQ(pb_result.indexed_vector_fields(1), "vec_field2");
}
TEST(ConverterTest, SegmentMetaWithEmptyFields) {
// Test conversion with minimal data
proto::SegmentMeta segment_pb;
segment_pb.set_segment_id(1);
auto segment_meta = ProtoConverter::FromPb(segment_pb);
ASSERT_NE(segment_meta, nullptr);
EXPECT_EQ(segment_meta->id(), 1u);
EXPECT_EQ(segment_meta->persisted_blocks().size(), 0);
EXPECT_FALSE(segment_meta->has_writing_forward_block());
EXPECT_EQ(segment_meta->indexed_vector_fields().size(), 0);
// Test conversion from C++ to protobuf with minimal data
SegmentMeta original_meta(5);
auto pb_result = ProtoConverter::ToPb(original_meta);
EXPECT_EQ(pb_result.segment_id(), 5u);
EXPECT_EQ(pb_result.persisted_blocks_size(), 0);
EXPECT_FALSE(pb_result.has_writing_forward_block());
EXPECT_EQ(pb_result.indexed_vector_fields_size(), 0);
}
// ==================== enable_rotate roundtrip tests ====================
TEST(ConverterTest, HnswIndexParamsWithEnableRotate) {
// C++ -> PB -> C++ roundtrip with enable_rotate = true
HnswIndexParams original(MetricType::COSINE, 16, 200, QuantizeType::INT8,
false, QuantizerParam(true));
EXPECT_TRUE(original.quantizer_param().enable_rotate());
auto pb = ProtoConverter::ToPb(&original);
EXPECT_TRUE(pb.base().quantizer_param().enable_rotate());
auto restored = ProtoConverter::FromPb(pb);
ASSERT_NE(restored, nullptr);
EXPECT_TRUE(restored->quantizer_param().enable_rotate());
EXPECT_TRUE(restored->enable_rotate()); // convenience getter
EXPECT_EQ(restored->metric_type(), MetricType::COSINE);
EXPECT_EQ(restored->m(), 16);
EXPECT_EQ(restored->ef_construction(), 200);
EXPECT_EQ(restored->quantize_type(), QuantizeType::INT8);
// C++ -> PB -> C++ roundtrip with enable_rotate = false
HnswIndexParams original_no_rot(MetricType::L2, 32, 100, QuantizeType::FP16);
auto pb2 = ProtoConverter::ToPb(&original_no_rot);
EXPECT_FALSE(pb2.base().quantizer_param().enable_rotate());
auto restored2 = ProtoConverter::FromPb(pb2);
ASSERT_NE(restored2, nullptr);
EXPECT_FALSE(restored2->quantizer_param().enable_rotate());
}
TEST(ConverterTest, FlatIndexParamsWithEnableRotate) {
FlatIndexParams original(MetricType::IP, QuantizeType::INT8,
QuantizerParam(true));
EXPECT_TRUE(original.quantizer_param().enable_rotate());
auto pb = ProtoConverter::ToPb(&original);
EXPECT_TRUE(pb.base().quantizer_param().enable_rotate());
auto restored = ProtoConverter::FromPb(pb);
ASSERT_NE(restored, nullptr);
EXPECT_TRUE(restored->quantizer_param().enable_rotate());
EXPECT_EQ(restored->metric_type(), MetricType::IP);
EXPECT_EQ(restored->quantize_type(), QuantizeType::INT8);
// enable_rotate = false
FlatIndexParams original_no_rot(MetricType::L2, QuantizeType::FP16);
auto pb2 = ProtoConverter::ToPb(&original_no_rot);
EXPECT_FALSE(pb2.base().quantizer_param().enable_rotate());
auto restored2 = ProtoConverter::FromPb(pb2);
EXPECT_FALSE(restored2->quantizer_param().enable_rotate());
}
TEST(ConverterTest, IVFIndexParamsWithEnableRotate) {
IVFIndexParams original(MetricType::COSINE, 256, 20, true, QuantizeType::INT8,
QuantizerParam(true));
EXPECT_TRUE(original.quantizer_param().enable_rotate());
auto pb = ProtoConverter::ToPb(&original);
EXPECT_TRUE(pb.base().quantizer_param().enable_rotate());
auto restored = ProtoConverter::FromPb(pb);
ASSERT_NE(restored, nullptr);
EXPECT_TRUE(restored->quantizer_param().enable_rotate());
EXPECT_EQ(restored->metric_type(), MetricType::COSINE);
EXPECT_EQ(restored->n_list(), 256);
EXPECT_EQ(restored->n_iters(), 20);
EXPECT_TRUE(restored->use_soar());
EXPECT_EQ(restored->quantize_type(), QuantizeType::INT8);
// enable_rotate = false
IVFIndexParams original_no_rot(MetricType::L2, 128, 10, false,
QuantizeType::FP16);
auto pb2 = ProtoConverter::ToPb(&original_no_rot);
EXPECT_FALSE(pb2.base().quantizer_param().enable_rotate());
auto restored2 = ProtoConverter::FromPb(pb2);
EXPECT_FALSE(restored2->quantizer_param().enable_rotate());
}
@@ -0,0 +1,291 @@
// 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 <gtest/gtest.h>
#include "db/index/common/type_helper.h"
using namespace zvec;
TEST(IndexTypeCodeBookTest, ProtoToCppConversion) {
// Test conversion from protobuf to C++ IndexType
EXPECT_EQ(IndexTypeCodeBook::Get(proto::IT_HNSW), IndexType::HNSW);
EXPECT_EQ(IndexTypeCodeBook::Get(proto::IT_FLAT), IndexType::FLAT);
EXPECT_EQ(IndexTypeCodeBook::Get(proto::IT_IVF), IndexType::IVF);
EXPECT_EQ(IndexTypeCodeBook::Get(proto::IT_INVERT), IndexType::INVERT);
EXPECT_EQ(IndexTypeCodeBook::Get(proto::IT_UNDEFINED), IndexType::UNDEFINED);
EXPECT_EQ(IndexTypeCodeBook::Get(static_cast<proto::IndexType>(999)),
IndexType::UNDEFINED);
}
TEST(IndexTypeCodeBookTest, CppToProtoConversion) {
// Test conversion from C++ IndexType to protobuf IndexType
EXPECT_EQ(IndexTypeCodeBook::Get(IndexType::HNSW), proto::IT_HNSW);
EXPECT_EQ(IndexTypeCodeBook::Get(IndexType::FLAT), proto::IT_FLAT);
EXPECT_EQ(IndexTypeCodeBook::Get(IndexType::IVF), proto::IT_IVF);
EXPECT_EQ(IndexTypeCodeBook::Get(IndexType::INVERT), proto::IT_INVERT);
EXPECT_EQ(IndexTypeCodeBook::Get(IndexType::UNDEFINED), proto::IT_UNDEFINED);
EXPECT_EQ(IndexTypeCodeBook::Get(static_cast<IndexType>(999)),
proto::IT_UNDEFINED);
}
TEST(IndexTypeCodeBookTest, CppToStringConversion) {
// Test conversion from C++ IndexType to string
EXPECT_EQ(IndexTypeCodeBook::AsString(IndexType::HNSW), "HNSW");
EXPECT_EQ(IndexTypeCodeBook::AsString(IndexType::INVERT), "INVERT");
EXPECT_EQ(IndexTypeCodeBook::AsString(IndexType::UNDEFINED), "UNDEFINED");
EXPECT_EQ(IndexTypeCodeBook::AsString(static_cast<IndexType>(999)),
"UNDEFINED");
}
TEST(DataTypeCodeBookTest, IsArrayType) {
// Test array type detection
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_BINARY));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_STRING));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_BOOL));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_INT32));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_INT64));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_UINT32));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_UINT64));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_FLOAT));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_DOUBLE));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_VECTOR_BINARY32));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_VECTOR_BINARY64));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_VECTOR_FP16));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_VECTOR_FP32));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_VECTOR_FP64));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_VECTOR_INT4));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_VECTOR_INT8));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_VECTOR_INT16));
EXPECT_FALSE(DataTypeCodeBook::IsArrayType(proto::DT_SPARSE_VECTOR_FP32));
EXPECT_TRUE(DataTypeCodeBook::IsArrayType(proto::DT_ARRAY_BINARY));
EXPECT_TRUE(DataTypeCodeBook::IsArrayType(proto::DT_ARRAY_STRING));
EXPECT_TRUE(DataTypeCodeBook::IsArrayType(proto::DT_ARRAY_BOOL));
EXPECT_TRUE(DataTypeCodeBook::IsArrayType(proto::DT_ARRAY_INT32));
EXPECT_TRUE(DataTypeCodeBook::IsArrayType(proto::DT_ARRAY_INT64));
EXPECT_TRUE(DataTypeCodeBook::IsArrayType(proto::DT_ARRAY_UINT32));
EXPECT_TRUE(DataTypeCodeBook::IsArrayType(proto::DT_ARRAY_UINT64));
EXPECT_TRUE(DataTypeCodeBook::IsArrayType(proto::DT_ARRAY_FLOAT));
EXPECT_TRUE(DataTypeCodeBook::IsArrayType(proto::DT_ARRAY_DOUBLE));
}
TEST(DataTypeCodeBookTest, ProtoToCppConversion) {
// Test conversion from protobuf to C++ DataType
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_BINARY), DataType::BINARY);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_STRING), DataType::STRING);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_BOOL), DataType::BOOL);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_INT32), DataType::INT32);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_INT64), DataType::INT64);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_UINT32), DataType::UINT32);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_UINT64), DataType::UINT64);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_FLOAT), DataType::FLOAT);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_DOUBLE), DataType::DOUBLE);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_VECTOR_BINARY32),
DataType::VECTOR_BINARY32);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_VECTOR_BINARY64),
DataType::VECTOR_BINARY64);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_VECTOR_FP16),
DataType::VECTOR_FP16);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_VECTOR_FP32),
DataType::VECTOR_FP32);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_VECTOR_FP64),
DataType::VECTOR_FP64);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_VECTOR_INT4),
DataType::VECTOR_INT4);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_VECTOR_INT8),
DataType::VECTOR_INT8);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_VECTOR_INT16),
DataType::VECTOR_INT16);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_SPARSE_VECTOR_FP32),
DataType::SPARSE_VECTOR_FP32);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_ARRAY_BINARY),
DataType::ARRAY_BINARY);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_ARRAY_STRING),
DataType::ARRAY_STRING);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_ARRAY_BOOL), DataType::ARRAY_BOOL);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_ARRAY_INT32),
DataType::ARRAY_INT32);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_ARRAY_INT64),
DataType::ARRAY_INT64);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_ARRAY_UINT32),
DataType::ARRAY_UINT32);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_ARRAY_UINT64),
DataType::ARRAY_UINT64);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_ARRAY_FLOAT),
DataType::ARRAY_FLOAT);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_ARRAY_DOUBLE),
DataType::ARRAY_DOUBLE);
EXPECT_EQ(DataTypeCodeBook::Get(proto::DT_UNDEFINED), DataType::UNDEFINED);
EXPECT_EQ(DataTypeCodeBook::Get(static_cast<proto::DataType>(999)),
DataType::UNDEFINED);
}
TEST(DataTypeCodeBookTest, CppToProtoConversion) {
// Test conversion from C++ DataType to protobuf DataType
EXPECT_EQ(DataTypeCodeBook::Get(DataType::BINARY), proto::DT_BINARY);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::STRING), proto::DT_STRING);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::BOOL), proto::DT_BOOL);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::INT32), proto::DT_INT32);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::INT64), proto::DT_INT64);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::UINT32), proto::DT_UINT32);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::UINT64), proto::DT_UINT64);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::FLOAT), proto::DT_FLOAT);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::DOUBLE), proto::DT_DOUBLE);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::VECTOR_BINARY32),
proto::DT_VECTOR_BINARY32);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::VECTOR_BINARY64),
proto::DT_VECTOR_BINARY64);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::VECTOR_FP16),
proto::DT_VECTOR_FP16);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::VECTOR_FP32),
proto::DT_VECTOR_FP32);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::VECTOR_FP64),
proto::DT_VECTOR_FP64);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::VECTOR_INT4),
proto::DT_VECTOR_INT4);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::VECTOR_INT8),
proto::DT_VECTOR_INT8);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::VECTOR_INT16),
proto::DT_VECTOR_INT16);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::SPARSE_VECTOR_FP16),
proto::DT_SPARSE_VECTOR_FP16);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::SPARSE_VECTOR_FP32),
proto::DT_SPARSE_VECTOR_FP32);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::ARRAY_BINARY),
proto::DT_ARRAY_BINARY);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::ARRAY_STRING),
proto::DT_ARRAY_STRING);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::ARRAY_BOOL), proto::DT_ARRAY_BOOL);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::ARRAY_INT32),
proto::DT_ARRAY_INT32);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::ARRAY_INT64),
proto::DT_ARRAY_INT64);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::ARRAY_UINT32),
proto::DT_ARRAY_UINT32);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::ARRAY_UINT64),
proto::DT_ARRAY_UINT64);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::ARRAY_FLOAT),
proto::DT_ARRAY_FLOAT);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::ARRAY_DOUBLE),
proto::DT_ARRAY_DOUBLE);
EXPECT_EQ(DataTypeCodeBook::Get(DataType::UNDEFINED), proto::DT_UNDEFINED);
EXPECT_EQ(DataTypeCodeBook::Get(static_cast<DataType>(999)),
proto::DT_UNDEFINED);
}
TEST(DataTypeCodeBookTest, CppToStringConversion) {
// Test conversion from C++ DataType to string
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::BINARY), "BINARY");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::STRING), "STRING");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::BOOL), "BOOL");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::INT32), "INT32");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::INT64), "INT64");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::UINT32), "UINT32");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::UINT64), "UINT64");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::FLOAT), "FLOAT");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::DOUBLE), "DOUBLE");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::VECTOR_BINARY32),
"VECTOR_BINARY32");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::VECTOR_BINARY64),
"VECTOR_BINARY64");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::VECTOR_FP16), "VECTOR_FP16");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::VECTOR_FP32), "VECTOR_FP32");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::VECTOR_FP64), "VECTOR_FP64");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::VECTOR_INT4), "VECTOR_INT4");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::VECTOR_INT8), "VECTOR_INT8");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::VECTOR_INT16), "VECTOR_INT16");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::ARRAY_BINARY), "ARRAY_BINARY");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::ARRAY_STRING), "ARRAY_STRING");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::ARRAY_BOOL), "ARRAY_BOOL");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::ARRAY_INT32), "ARRAY_INT32");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::ARRAY_INT64), "ARRAY_INT64");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::ARRAY_UINT32), "ARRAY_UINT32");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::ARRAY_UINT64), "ARRAY_UINT64");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::ARRAY_FLOAT), "ARRAY_FLOAT");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::ARRAY_DOUBLE), "ARRAY_DOUBLE");
EXPECT_EQ(DataTypeCodeBook::AsString(DataType::UNDEFINED), "");
EXPECT_EQ(DataTypeCodeBook::AsString(static_cast<DataType>(999)), "");
}
TEST(MetricTypeCodeBookTest, ProtoToCppConversion) {
// Test conversion from protobuf to C++ MetricType
EXPECT_EQ(MetricTypeCodeBook::Get(proto::MT_IP), MetricType::IP);
EXPECT_EQ(MetricTypeCodeBook::Get(proto::MT_L2), MetricType::L2);
EXPECT_EQ(MetricTypeCodeBook::Get(proto::MT_COSINE), MetricType::COSINE);
EXPECT_EQ(MetricTypeCodeBook::Get(proto::MT_UNDEFINED),
MetricType::UNDEFINED);
EXPECT_EQ(MetricTypeCodeBook::Get(static_cast<proto::MetricType>(999)),
MetricType::UNDEFINED);
}
TEST(MetricTypeCodeBookTest, CppToProtoConversion) {
// Test conversion from C++ MetricType to protobuf MetricType
EXPECT_EQ(MetricTypeCodeBook::Get(MetricType::IP), proto::MT_IP);
EXPECT_EQ(MetricTypeCodeBook::Get(MetricType::L2), proto::MT_L2);
EXPECT_EQ(MetricTypeCodeBook::Get(MetricType::COSINE), proto::MT_COSINE);
EXPECT_EQ(MetricTypeCodeBook::Get(MetricType::UNDEFINED),
proto::MT_UNDEFINED);
EXPECT_EQ(MetricTypeCodeBook::Get(static_cast<MetricType>(999)),
proto::MT_UNDEFINED);
}
TEST(QuantizeTypeCodeBookTest, ProtoToCppConversion) {
// Test conversion from protobuf to C++ QuantizeType
EXPECT_EQ(QuantizeTypeCodeBook::Get(proto::QT_FP16), QuantizeType::FP16);
EXPECT_EQ(QuantizeTypeCodeBook::Get(proto::QT_INT4), QuantizeType::INT4);
EXPECT_EQ(QuantizeTypeCodeBook::Get(proto::QT_INT8), QuantizeType::INT8);
EXPECT_EQ(QuantizeTypeCodeBook::Get(proto::QT_UNDEFINED),
QuantizeType::UNDEFINED);
EXPECT_EQ(QuantizeTypeCodeBook::Get(static_cast<proto::QuantizeType>(999)),
QuantizeType::UNDEFINED);
}
TEST(QuantizeTypeCodeBookTest, CppToProtoConversion) {
// Test conversion from C++ QuantizeType to protobuf QuantizeType
EXPECT_EQ(QuantizeTypeCodeBook::Get(QuantizeType::FP16), proto::QT_FP16);
EXPECT_EQ(QuantizeTypeCodeBook::Get(QuantizeType::INT4), proto::QT_INT4);
EXPECT_EQ(QuantizeTypeCodeBook::Get(QuantizeType::INT8), proto::QT_INT8);
EXPECT_EQ(QuantizeTypeCodeBook::Get(QuantizeType::UNDEFINED),
proto::QT_UNDEFINED);
EXPECT_EQ(QuantizeTypeCodeBook::Get(static_cast<QuantizeType>(999)),
proto::QT_UNDEFINED);
}
TEST(BlockTypeCodeBookTest, ProtoToCppConversion) {
// Test conversion from protobuf to C++ BlockType
EXPECT_EQ(BlockTypeCodeBook::Get(proto::BT_SCALAR), BlockType::SCALAR);
EXPECT_EQ(BlockTypeCodeBook::Get(proto::BT_SCALAR_INDEX),
BlockType::SCALAR_INDEX);
EXPECT_EQ(BlockTypeCodeBook::Get(proto::BT_VECTOR_INDEX),
BlockType::VECTOR_INDEX);
EXPECT_EQ(BlockTypeCodeBook::Get(proto::BT_VECTOR_INDEX_QUANTIZE),
BlockType::VECTOR_INDEX_QUANTIZE);
EXPECT_EQ(BlockTypeCodeBook::Get(proto::BT_UNDEFINED), BlockType::UNDEFINED);
EXPECT_EQ(BlockTypeCodeBook::Get(static_cast<proto::BlockType>(999)),
BlockType::UNDEFINED);
}
TEST(BlockTypeCodeBookTest, CppToProtoConversion) {
// Test conversion from C++ BlockType to protobuf BlockType
EXPECT_EQ(BlockTypeCodeBook::Get(BlockType::SCALAR), proto::BT_SCALAR);
EXPECT_EQ(BlockTypeCodeBook::Get(BlockType::SCALAR_INDEX),
proto::BT_SCALAR_INDEX);
EXPECT_EQ(BlockTypeCodeBook::Get(BlockType::VECTOR_INDEX),
proto::BT_VECTOR_INDEX);
EXPECT_EQ(BlockTypeCodeBook::Get(BlockType::VECTOR_INDEX_QUANTIZE),
proto::BT_VECTOR_INDEX_QUANTIZE);
EXPECT_EQ(BlockTypeCodeBook::Get(BlockType::UNDEFINED), proto::BT_UNDEFINED);
EXPECT_EQ(BlockTypeCodeBook::Get(static_cast<BlockType>(999)),
proto::BT_UNDEFINED);
}
File diff suppressed because it is too large Load Diff
+281
View File
@@ -0,0 +1,281 @@
// 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(), &params); // 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(), &params);
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(), &params);
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(), &params);
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);
}
}
+252
View File
@@ -0,0 +1,252 @@
// 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 "db/index/common/meta.h"
#include <gtest/gtest.h>
using namespace zvec;
TEST(SegmentMetaTest, DefaultConstruction) {
SegmentMeta segment_meta;
EXPECT_EQ(segment_meta.id(), 0u);
EXPECT_TRUE(segment_meta.persisted_blocks().empty());
EXPECT_FALSE(segment_meta.has_writing_forward_block());
EXPECT_EQ(segment_meta.min_doc_id(), 0u);
EXPECT_EQ(segment_meta.max_doc_id(), 0u);
EXPECT_EQ(segment_meta.doc_count(), 0u);
EXPECT_TRUE(segment_meta.indexed_vector_fields().empty());
}
TEST(SegmentMetaTest, ConstructionWithID) {
SegmentMeta segment_meta(42);
EXPECT_EQ(segment_meta.id(), 42u);
EXPECT_TRUE(segment_meta.persisted_blocks().empty());
EXPECT_FALSE(segment_meta.has_writing_forward_block());
EXPECT_EQ(segment_meta.min_doc_id(), 0u);
EXPECT_EQ(segment_meta.max_doc_id(), 0u);
EXPECT_EQ(segment_meta.doc_count(), 0u);
EXPECT_TRUE(segment_meta.indexed_vector_fields().empty());
}
TEST(SegmentMetaTest, PersistedBlocksOperations) {
SegmentMeta segment_meta(1);
// Add persisted blocks
BlockMeta block1(1, BlockType::SCALAR, 0, 100);
block1.set_doc_count(50);
block1.add_column("col1");
block1.add_column("col2");
BlockMeta block2(2, BlockType::VECTOR_INDEX, 101, 200);
block2.set_doc_count(75);
block2.add_column("vec_col");
segment_meta.add_persisted_block(block1);
segment_meta.add_persisted_block(block2);
EXPECT_EQ(segment_meta.persisted_blocks().size(), 2u);
const auto &blocks = segment_meta.persisted_blocks();
EXPECT_EQ(blocks[0].id(), 1u);
EXPECT_EQ(blocks[0].type(), BlockType::SCALAR);
EXPECT_EQ(blocks[0].min_doc_id(), 0u);
EXPECT_EQ(blocks[0].max_doc_id(), 100u);
EXPECT_EQ(blocks[0].doc_count(), 50u);
EXPECT_EQ(blocks[0].columns().size(), 2u);
EXPECT_EQ(blocks[1].id(), 2u);
EXPECT_EQ(blocks[1].type(), BlockType::VECTOR_INDEX);
EXPECT_EQ(blocks[1].min_doc_id(), 101u);
EXPECT_EQ(blocks[1].max_doc_id(), 200u);
EXPECT_EQ(blocks[1].doc_count(), 75u);
EXPECT_EQ(blocks[1].columns().size(), 1u);
}
TEST(SegmentMetaTest, WritingForwardBlockOperations) {
SegmentMeta segment_meta(1);
// Initially no writing forward block
EXPECT_FALSE(segment_meta.has_writing_forward_block());
// Set writing forward block
BlockMeta writing_block(3, BlockType::SCALAR, 201, 300);
writing_block.set_doc_count(25);
writing_block.add_column("col3");
segment_meta.set_writing_forward_block(writing_block);
// Now should have writing forward block
EXPECT_TRUE(segment_meta.has_writing_forward_block());
const auto &wfb = segment_meta.writing_forward_block();
EXPECT_EQ(wfb.value().id(), 3u);
EXPECT_EQ(wfb.value().type(), BlockType::SCALAR);
EXPECT_EQ(wfb.value().min_doc_id(), 201u);
EXPECT_EQ(wfb.value().max_doc_id(), 300u);
EXPECT_EQ(wfb.value().doc_count(), 25u);
EXPECT_EQ(wfb.value().columns().size(), 1u);
EXPECT_EQ(wfb.value().columns()[0], "col3");
}
TEST(SegmentMetaTest, MinDocIDCalculation) {
SegmentMeta segment_meta(1);
// Case 1: No persisted blocks, no writing forward block
EXPECT_EQ(segment_meta.min_doc_id(), 0u);
// Case 2: No persisted blocks, but has writing forward block
BlockMeta writing_block(1, BlockType::SCALAR, 100, 200);
segment_meta.set_writing_forward_block(writing_block);
EXPECT_EQ(segment_meta.min_doc_id(), 100u);
// Case 3: Has persisted blocks (should take precedence)
BlockMeta persisted_block(1, BlockType::SCALAR, 50, 150);
segment_meta.add_persisted_block(persisted_block);
EXPECT_EQ(segment_meta.min_doc_id(), 50u);
}
TEST(SegmentMetaTest, MaxDocIDCalculation) {
SegmentMeta segment_meta(1);
// Case 1: No blocks at all
EXPECT_EQ(segment_meta.max_doc_id(), 0u);
// Case 2: Only persisted blocks
BlockMeta persisted_block(1, BlockType::SCALAR, 0, 100);
segment_meta.add_persisted_block(persisted_block);
EXPECT_EQ(segment_meta.max_doc_id(), 100u);
// Case 3: Both persisted and writing forward blocks (writing forward takes
// precedence)
BlockMeta writing_block(2, BlockType::SCALAR, 101, 200);
segment_meta.set_writing_forward_block(writing_block);
EXPECT_EQ(segment_meta.max_doc_id(), 100u);
// Case 4: Only writing forward block
SegmentMeta segment_meta2(2);
segment_meta2.set_writing_forward_block(writing_block);
EXPECT_EQ(segment_meta2.max_doc_id(), 0u);
}
TEST(SegmentMetaTest, DocCountCalculation) {
SegmentMeta segment_meta(1);
// Initially 0
EXPECT_EQ(segment_meta.doc_count(), 0u);
// Add persisted blocks
BlockMeta block1(1, BlockType::SCALAR, 0, 100);
block1.set_doc_count(50);
segment_meta.add_persisted_block(block1);
EXPECT_EQ(segment_meta.doc_count(), 50u);
// Add another persisted block
BlockMeta block2(2, BlockType::VECTOR_INDEX, 101, 200);
block2.set_doc_count(75);
segment_meta.add_persisted_block(block2);
EXPECT_EQ(segment_meta.doc_count(), 50u);
// Add writing forward block
BlockMeta writing_block(3, BlockType::SCALAR, 201, 300);
writing_block.set_doc_count(25);
segment_meta.set_writing_forward_block(writing_block);
EXPECT_EQ(segment_meta.doc_count(), 75);
}
TEST(SegmentMetaTest, IndexedVectorFieldsOperations) {
SegmentMeta segment_meta(1);
// Initially empty
EXPECT_FALSE(segment_meta.vector_indexed("field1"));
EXPECT_TRUE(segment_meta.indexed_vector_fields().empty());
// Add indexed fields
segment_meta.add_indexed_vector_field("field1");
segment_meta.add_indexed_vector_field("field2");
EXPECT_TRUE(segment_meta.vector_indexed("field1"));
EXPECT_TRUE(segment_meta.vector_indexed("field2"));
EXPECT_FALSE(segment_meta.vector_indexed("field3"));
EXPECT_EQ(segment_meta.indexed_vector_fields().size(), 2u);
// Check set operation
std::set<std::string> fields = {"field3", "field4"};
segment_meta.set_indexed_vector_fields(fields);
EXPECT_FALSE(segment_meta.vector_indexed("field1"));
EXPECT_FALSE(segment_meta.vector_indexed("field2"));
EXPECT_TRUE(segment_meta.vector_indexed("field3"));
EXPECT_TRUE(segment_meta.vector_indexed("field4"));
EXPECT_EQ(segment_meta.indexed_vector_fields().size(), 2u);
}
TEST(SegmentMetaTest, UpdateMaxDocId) {
SegmentMeta segment_meta(1);
// Try to update when no writing forward block - should not crash
segment_meta.update_max_doc_id(100);
// Set writing forward block and update
BlockMeta writing_block(1, BlockType::SCALAR, 0, 50);
segment_meta.set_writing_forward_block(writing_block);
EXPECT_EQ(segment_meta.writing_forward_block().value().max_doc_id(), 50u);
segment_meta.update_max_doc_id(100);
EXPECT_EQ(segment_meta.writing_forward_block().value().max_doc_id(), 100u);
}
TEST(SegmentMetaTest, EqualityOperators) {
SegmentMeta segment1(1);
SegmentMeta segment2(1);
SegmentMeta segment3(2);
// Same empty segments
EXPECT_TRUE(segment1 == segment2);
EXPECT_FALSE(segment1 != segment2);
// Different IDs
EXPECT_FALSE(segment1 == segment3);
EXPECT_TRUE(segment1 != segment3);
// Add same persisted block to both
BlockMeta block(1, BlockType::SCALAR, 0, 100);
block.set_doc_count(50);
segment1.add_persisted_block(block);
segment2.add_persisted_block(block);
EXPECT_TRUE(segment1 == segment2);
// Add writing forward block
BlockMeta wfb(2, BlockType::VECTOR_INDEX, 101, 200);
segment1.set_writing_forward_block(wfb);
segment2.set_writing_forward_block(wfb);
EXPECT_TRUE(segment1 == segment2);
// Add indexed fields
segment1.add_indexed_vector_field("vec_field");
segment2.add_indexed_vector_field("vec_field");
EXPECT_TRUE(segment1 == segment2);
// Make them different again
segment1.add_indexed_vector_field("vec_field2");
EXPECT_FALSE(segment1 == segment2);
EXPECT_TRUE(segment1 != segment2);
}
@@ -0,0 +1,80 @@
// 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/query_params.h"
#include <gtest/gtest.h>
using namespace zvec;
TEST(QueryParamsTest, QueryParamsBaseClass) {
// Test constructor
QueryParams params(IndexType::HNSW);
EXPECT_EQ(params.type(), IndexType::HNSW);
}
TEST(QueryParamsTest, HnswQueryParams) {
// Test constructor
HnswQueryParams params(100);
EXPECT_EQ(params.type(), IndexType::HNSW);
EXPECT_EQ(params.ef(), 100);
// Test setter
params.set_ef(200);
EXPECT_EQ(params.ef(), 200);
}
TEST(QueryParamsTest, IVFQueryParams) {
// Test constructor
IVFQueryParams params(50);
EXPECT_EQ(params.type(), IndexType::IVF);
EXPECT_EQ(params.nprobe(), 50);
// Test setter
params.set_nprobe(75);
EXPECT_EQ(params.nprobe(), 75);
}
TEST(QueryParamsTest, Polymorphism) {
// Test polymorphic behavior
QueryParams::Ptr hnsw_ptr = std::make_shared<HnswQueryParams>(100);
QueryParams::Ptr ivf_ptr = std::make_shared<IVFQueryParams>(50);
// Verify types
EXPECT_EQ(hnsw_ptr->type(), IndexType::HNSW);
EXPECT_EQ(ivf_ptr->type(), IndexType::IVF);
// Test dynamic casting
auto hnsw_cast = std::dynamic_pointer_cast<HnswQueryParams>(hnsw_ptr);
auto ivf_cast = std::dynamic_pointer_cast<IVFQueryParams>(ivf_ptr);
auto invalid_cast = std::dynamic_pointer_cast<HnswQueryParams>(ivf_ptr);
EXPECT_NE(hnsw_cast, nullptr);
EXPECT_NE(ivf_cast, nullptr);
EXPECT_EQ(invalid_cast, nullptr);
// Verify values after casting
EXPECT_EQ(hnsw_cast->ef(), 100);
EXPECT_EQ(ivf_cast->nprobe(), 50);
}
TEST(QueryParamsTest, VirtualDestructor) {
// Test that virtual destructor allows proper deletion
QueryParams *hnsw_ptr = new HnswQueryParams(100);
QueryParams *ivf_ptr = new IVFQueryParams(50);
// This should not cause memory issues
delete hnsw_ptr;
delete ivf_ptr;
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,303 @@
// 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 "db/index/common/version_manager.h"
#include <filesystem>
#include <memory>
#include <gtest/gtest.h>
#include "db/common/file_helper.h"
#include "db/index/common/meta.h"
#include "proto/zvec.pb.h"
#include "zvec/db/schema.h"
namespace zvec {
class VersionManagerTest : public ::testing::Test {
protected:
void SetUp() override {
// Create a temporary directory for testing
test_path_ = "./test_version_manager";
FileHelper::RemoveDirectory(test_path_);
ASSERT_TRUE(FileHelper::CreateDirectory(test_path_))
<< ailego::FileHelper::GetLastErrorString();
}
void TearDown() override {
// Clean up temporary files
FileHelper::RemoveDirectory(test_path_);
}
std::string test_path_;
};
// Test basic Version functionality
TEST_F(VersionManagerTest, VersionBasicOperations) {
Version version;
// Create a sample schema
CollectionSchema schema;
schema.set_name("test_collection");
version.set_schema(schema);
// Verify schema is set correctly
EXPECT_EQ(version.schema().name(), "test_collection");
// Test segment meta operations
auto segment_meta = std::make_shared<SegmentMeta>(1);
segment_meta->set_id(1);
// Add segment meta
EXPECT_TRUE(version.add_persisted_segment_meta(segment_meta).ok());
// Try to add duplicate - should fail
EXPECT_FALSE(version.add_persisted_segment_meta(segment_meta).ok());
// Get segment metas
auto segment_metas = version.persisted_segment_metas();
EXPECT_EQ(segment_metas.size(), 1);
EXPECT_EQ(segment_metas[0]->id(), 1);
// Remove segment meta
EXPECT_TRUE(version.remove_persisted_segment_meta(1).ok());
EXPECT_EQ(version.persisted_segment_metas().size(), 0);
// Try to remove non-existent segment - should fail
EXPECT_FALSE(version.remove_persisted_segment_meta(1).ok());
std::cout << version.to_string() << std::endl;
std::cout << version.to_string_formatted() << std::endl;
}
// Test Version Load/Save operations
TEST_F(VersionManagerTest, VersionLoadSave) {
std::string manifest_path = test_path_ + "/manifest";
// Create and populate a version
Version version;
CollectionSchema schema;
schema.set_name("test_collection");
version.set_schema(schema);
auto segment_meta = std::make_shared<SegmentMeta>(1);
segment_meta->set_id(1);
version.add_persisted_segment_meta(segment_meta);
version.set_id_map_path_suffix(100);
version.set_delete_snapshot_path_suffix(200);
version.set_next_segment_id(2);
// Save version
EXPECT_TRUE(Version::Save(manifest_path, version).ok());
// Load version
Version loaded_version;
EXPECT_TRUE(Version::Load(manifest_path, &loaded_version).ok());
// Verify loaded version matches original
EXPECT_EQ(loaded_version.schema().name(), "test_collection");
EXPECT_EQ(loaded_version.persisted_segment_metas().size(), 1);
EXPECT_EQ(loaded_version.id_map_path_suffix(), 100);
EXPECT_EQ(loaded_version.delete_snapshot_path_suffix(), 200);
EXPECT_EQ(loaded_version.next_segment_id(), 2);
}
// Test VersionManager creation and recovery
TEST_F(VersionManagerTest, VersionManagerCreateAndRecover) {
std::string version_path = test_path_ + "/version";
std::filesystem::create_directories(version_path);
// Create initial version
Version initial_version;
CollectionSchema schema;
schema.set_name("initial_collection");
initial_version.set_schema(schema);
auto segment_meta = std::make_shared<SegmentMeta>(1);
segment_meta->set_id(1);
initial_version.add_persisted_segment_meta(segment_meta);
// Create VersionManager
auto create_result = VersionManager::Create(version_path, initial_version);
ASSERT_TRUE(create_result.has_value());
auto version_manager = create_result.value();
// Get current version and verify
auto current_version = version_manager->get_current_version();
EXPECT_EQ(current_version.schema().name(), "initial_collection");
// Modify version
auto new_segment = std::make_shared<SegmentMeta>(2);
new_segment->set_id(2);
EXPECT_TRUE(version_manager->add_persisted_segment_meta(new_segment).ok());
// Flush changes
ASSERT_TRUE(version_manager->flush().ok());
// Recover VersionManager
auto recover_result = VersionManager::Recovery(version_path);
ASSERT_TRUE(recover_result.has_value());
auto recovered_manager = recover_result.value();
auto recovered_version = recovered_manager->get_current_version();
// Verify recovered version matches modified version
EXPECT_EQ(recovered_version.schema().name(), "initial_collection");
EXPECT_EQ(recovered_version.persisted_segment_metas().size(), 2);
}
// Test VersionManager operations
TEST_F(VersionManagerTest, VersionManagerOperations) {
std::string version_path = test_path_ + "/version_ops";
std::filesystem::create_directories(version_path);
// Create initial version
Version initial_version;
CollectionSchema schema;
schema.set_name("test_collection");
initial_version.set_schema(schema);
auto create_result = VersionManager::Create(version_path, initial_version);
auto version_manager = create_result.value();
// Test segment meta operations through VersionManager
auto segment_meta = std::make_shared<SegmentMeta>(1);
segment_meta->set_id(1);
EXPECT_TRUE(version_manager->add_persisted_segment_meta(segment_meta).ok());
// Test reset writing segment meta
auto writing_segment = std::make_shared<SegmentMeta>(100);
writing_segment->set_id(100);
EXPECT_TRUE(
version_manager->reset_writing_segment_meta(writing_segment).ok());
// Test suffix setters
version_manager->set_id_map_path_suffix(50);
version_manager->set_delete_snapshot_path_suffix(60);
version_manager->set_next_segment_id(3);
// Flush and verify
EXPECT_TRUE(version_manager->flush().ok());
auto current_version = version_manager->get_current_version();
EXPECT_EQ(current_version.id_map_path_suffix(), 50);
EXPECT_EQ(current_version.delete_snapshot_path_suffix(), 60);
EXPECT_EQ(current_version.next_segment_id(), 3);
EXPECT_EQ(current_version.writing_segment_meta()->id(), 100);
}
// Test Version equality operator
TEST_F(VersionManagerTest, VersionEquality) {
Version version1, version2;
CollectionSchema schema1, schema2;
schema1.set_name("collection1");
schema2.set_name("collection1");
version1.set_schema(schema1);
version2.set_schema(schema2);
auto segment_meta1 = std::make_shared<SegmentMeta>(1);
segment_meta1->set_id(1);
version1.add_persisted_segment_meta(segment_meta1);
auto segment_meta2 = std::make_shared<SegmentMeta>(1);
segment_meta2->set_id(1);
version2.add_persisted_segment_meta(segment_meta2);
// Versions should be equal
EXPECT_TRUE(version1 == version2);
// Make them different
auto segment_meta3 = std::make_shared<SegmentMeta>(2);
segment_meta3->set_id(2);
version2.add_persisted_segment_meta(segment_meta3);
// Versions should not be equal now
EXPECT_FALSE(version1 == version2);
}
// Test error conditions
TEST_F(VersionManagerTest, ErrorConditions) {
std::string version_path = test_path_ + "/error_test";
std::filesystem::create_directories(version_path);
// Create initial version
Version initial_version;
CollectionSchema schema;
schema.set_name("test");
initial_version.set_schema(schema);
auto create_result = VersionManager::Create(version_path, initial_version);
auto version_manager = create_result.value();
// Test operations with null segment meta
EXPECT_FALSE(version_manager->add_persisted_segment_meta(nullptr).ok());
// Test operations with non-existent segment ID
EXPECT_FALSE(version_manager->remove_persisted_segment_meta(999).ok());
}
// Test conversion between protobuf and internal schema
TEST_F(VersionManagerTest, SchemaConversion) {
// Create protobuf schema
zvec::proto::CollectionSchema pb_schema;
pb_schema.set_name("test_collection");
auto pb_field = pb_schema.add_fields();
pb_field->set_name("vector_field");
pb_field->set_data_type(zvec::proto::DataType::DT_VECTOR_FP32);
pb_field->set_dimension(128);
// Convert to internal schema (this would be done in the Load method)
CollectionSchema internal_schema;
internal_schema.set_name(pb_schema.name());
// In a real implementation, fields would be converted here
// Test that we can set and retrieve the schema
Version version;
version.set_schema(internal_schema);
EXPECT_EQ(version.schema().name(), "test_collection");
}
// Test SegmentMeta functionality
TEST_F(VersionManagerTest, SegmentMetaOperations) {
SegmentMeta segment_meta(10);
EXPECT_EQ(segment_meta.id(), 10);
// Test block operations
BlockMeta block(1, BlockType::SCALAR, 0, 100);
segment_meta.add_persisted_block(block);
EXPECT_EQ(segment_meta.persisted_blocks().size(), 1);
EXPECT_EQ(segment_meta.persisted_blocks()[0].id(), 1);
// Test indexed vector fields
EXPECT_FALSE(segment_meta.vector_indexed("field1"));
segment_meta.add_indexed_vector_field("field1");
EXPECT_TRUE(segment_meta.vector_indexed("field1"));
// Test min/max doc id
EXPECT_EQ(segment_meta.min_doc_id(), 0);
EXPECT_EQ(segment_meta.max_doc_id(), 100);
}
} // namespace zvec
@@ -0,0 +1,437 @@
// 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 "db/index/segment/column_merging_reader.h"
#include <memory>
#include <vector>
#include <arrow/api.h>
#include <arrow/builder.h>
#include <arrow/ipc/writer.h>
#include <arrow/testing/gtest_util.h>
#include <arrow/testing/util.h>
#include <gtest/gtest.h>
using namespace zvec;
arrow::Result<std::shared_ptr<arrow::Array>> MakeInt32Array(
const std::vector<int32_t> &values) {
arrow::Int32Builder builder;
ARROW_RETURN_NOT_OK(builder.AppendValues(values));
std::shared_ptr<arrow::Array> array;
ARROW_RETURN_NOT_OK(builder.Finish(&array));
return array;
}
arrow::Result<std::shared_ptr<arrow::RecordBatch>> MakeInt32RecordBatch(
const std::string &column_name, const std::vector<int32_t> &values) {
ARROW_ASSIGN_OR_RAISE(auto array, MakeInt32Array(values));
auto schema = arrow::schema({arrow::field(column_name, arrow::int32())});
return arrow::RecordBatch::Make(schema, values.size(), {array});
}
// Mock RecordBatchReader for testing error conditions
class MockErrorRecordBatchReader : public arrow::ipc::RecordBatchReader {
public:
explicit MockErrorRecordBatchReader(arrow::StatusCode error_code)
: error_code_(error_code) {}
std::shared_ptr<arrow::Schema> schema() const override {
return arrow::schema({arrow::field("dummy", arrow::int32())});
}
arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch> *out) override {
*out = nullptr;
return arrow::Status(error_code_, "Mock error");
}
private:
arrow::StatusCode error_code_;
};
// Test fixture
class ColumnMergingReaderTest : public ::testing::Test {
protected:
void SetUp() override {
// Create test schemas
schema1_ = arrow::schema({arrow::field("col1", arrow::int32()),
arrow::field("col2", arrow::int32())});
schema2_ = arrow::schema({arrow::field("col3", arrow::int32()),
arrow::field("col4", arrow::int32())});
target_schema_ = arrow::schema({arrow::field("col1", arrow::int32()),
arrow::field("col2", arrow::int32()),
arrow::field("col3", arrow::int32()),
arrow::field("col4", arrow::int32())});
}
std::shared_ptr<arrow::Schema> schema1_;
std::shared_ptr<arrow::Schema> schema2_;
std::shared_ptr<arrow::Schema> target_schema_;
};
// Test Make factory method
TEST_F(ColumnMergingReaderTest, Make) {
std::vector<std::shared_ptr<arrow::ipc::RecordBatchReader>> readers;
auto reader = ColumnMergingReader::Make(target_schema_, std::move(readers));
ASSERT_NE(reader, nullptr);
EXPECT_EQ(reader->schema(), target_schema_);
}
// Test constructor and schema method
TEST_F(ColumnMergingReaderTest, ConstructorAndSchema) {
std::vector<std::shared_ptr<arrow::ipc::RecordBatchReader>> readers;
auto reader =
std::make_shared<ColumnMergingReader>(target_schema_, std::move(readers));
EXPECT_EQ(reader->schema(), target_schema_);
}
// Test normal operation with two readers
TEST_F(ColumnMergingReaderTest, NormalOperation) {
// Create first batch with col1 and col2
auto array1 = MakeInt32Array({1, 2, 3}).ValueOrDie();
auto array2 = MakeInt32Array({4, 5, 6}).ValueOrDie();
auto batch1 = arrow::RecordBatch::Make(schema1_, 3, {array1, array2});
// Create second batch with col3 and col4
auto array3 = MakeInt32Array({7, 8, 9}).ValueOrDie();
auto array4 = MakeInt32Array({10, 11, 12}).ValueOrDie();
auto batch2 = arrow::RecordBatch::Make(schema2_, 3, {array3, array4});
// Create mock readers
class MockRecordBatchReader : public arrow::ipc::RecordBatchReader {
public:
explicit MockRecordBatchReader(std::shared_ptr<arrow::RecordBatch> batch)
: batch_(batch), returned_(false) {}
std::shared_ptr<arrow::Schema> schema() const override {
return batch_->schema();
}
arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch> *out) override {
if (!returned_) {
*out = batch_;
returned_ = true;
} else {
*out = nullptr;
}
return arrow::Status::OK();
}
private:
std::shared_ptr<arrow::RecordBatch> batch_;
bool returned_;
};
std::vector<std::shared_ptr<arrow::ipc::RecordBatchReader>> readers;
readers.push_back(std::make_shared<MockRecordBatchReader>(batch1));
readers.push_back(std::make_shared<MockRecordBatchReader>(batch2));
auto merging_reader =
ColumnMergingReader::Make(target_schema_, std::move(readers));
std::shared_ptr<arrow::RecordBatch> result_batch;
ASSERT_OK(merging_reader->ReadNext(&result_batch));
ASSERT_NE(result_batch, nullptr);
EXPECT_EQ(result_batch->num_rows(), 3);
EXPECT_EQ(result_batch->num_columns(), 4);
// Check column values
auto col1 =
std::static_pointer_cast<arrow::Int32Array>(result_batch->column(0));
auto col2 =
std::static_pointer_cast<arrow::Int32Array>(result_batch->column(1));
auto col3 =
std::static_pointer_cast<arrow::Int32Array>(result_batch->column(2));
auto col4 =
std::static_pointer_cast<arrow::Int32Array>(result_batch->column(3));
EXPECT_EQ(col1->Value(0), 1);
EXPECT_EQ(col1->Value(1), 2);
EXPECT_EQ(col1->Value(2), 3);
EXPECT_EQ(col2->Value(0), 4);
EXPECT_EQ(col2->Value(1), 5);
EXPECT_EQ(col2->Value(2), 6);
EXPECT_EQ(col3->Value(0), 7);
EXPECT_EQ(col3->Value(1), 8);
EXPECT_EQ(col3->Value(2), 9);
EXPECT_EQ(col4->Value(0), 10);
EXPECT_EQ(col4->Value(1), 11);
EXPECT_EQ(col4->Value(2), 12);
// Second read should return nullptr (EOF)
ASSERT_OK(merging_reader->ReadNext(&result_batch));
EXPECT_EQ(result_batch, nullptr);
}
// Test with empty readers
TEST_F(ColumnMergingReaderTest, EmptyReaders) {
std::vector<std::shared_ptr<arrow::ipc::RecordBatchReader>> readers;
auto merging_reader =
ColumnMergingReader::Make(target_schema_, std::move(readers));
std::shared_ptr<arrow::RecordBatch> result_batch;
ASSERT_OK(merging_reader->ReadNext(&result_batch));
EXPECT_EQ(result_batch, nullptr);
}
// Test with inconsistent row counts
TEST_F(ColumnMergingReaderTest, InconsistentRowCounts) {
// Create first batch with 3 rows
auto array1 = MakeInt32Array({1, 2, 3}).ValueOrDie();
auto batch1 = arrow::RecordBatch::Make(
arrow::schema({arrow::field("col1", arrow::int32())}), 3, {array1});
// Create second batch with 2 rows
auto array2 = MakeInt32Array({4, 5}).ValueOrDie();
auto batch2 = arrow::RecordBatch::Make(
arrow::schema({arrow::field("col2", arrow::int32())}), 2, {array2});
class MockRecordBatchReader : public arrow::ipc::RecordBatchReader {
public:
explicit MockRecordBatchReader(std::shared_ptr<arrow::RecordBatch> batch)
: batch_(batch), returned_(false) {}
std::shared_ptr<arrow::Schema> schema() const override {
return batch_->schema();
}
arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch> *out) override {
if (!returned_) {
*out = batch_;
returned_ = true;
} else {
*out = nullptr;
}
return arrow::Status::OK();
}
private:
std::shared_ptr<arrow::RecordBatch> batch_;
bool returned_;
};
std::vector<std::shared_ptr<arrow::ipc::RecordBatchReader>> readers;
readers.push_back(std::make_shared<MockRecordBatchReader>(batch1));
readers.push_back(std::make_shared<MockRecordBatchReader>(batch2));
auto merging_reader =
ColumnMergingReader::Make(target_schema_, std::move(readers));
std::shared_ptr<arrow::RecordBatch> result_batch;
arrow::Status status = merging_reader->ReadNext(&result_batch);
EXPECT_FALSE(status.ok());
EXPECT_EQ(status.code(), arrow::StatusCode::Invalid);
}
// Test missing column
TEST_F(ColumnMergingReaderTest, MissingColumn) {
// Create batch with only col1
auto array1 = MakeInt32Array({1, 2, 3}).ValueOrDie();
auto batch1 = arrow::RecordBatch::Make(
arrow::schema({arrow::field("col1", arrow::int32())}), 3, {array1});
class MockRecordBatchReader : public arrow::ipc::RecordBatchReader {
public:
explicit MockRecordBatchReader(std::shared_ptr<arrow::RecordBatch> batch)
: batch_(batch), returned_(false) {}
std::shared_ptr<arrow::Schema> schema() const override {
return batch_->schema();
}
arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch> *out) override {
if (!returned_) {
*out = batch_;
returned_ = true;
} else {
*out = nullptr;
}
return arrow::Status::OK();
}
private:
std::shared_ptr<arrow::RecordBatch> batch_;
bool returned_;
};
std::vector<std::shared_ptr<arrow::ipc::RecordBatchReader>> readers;
readers.push_back(std::make_shared<MockRecordBatchReader>(batch1));
// Target schema requires col1 and col2 but we only provide col1
auto target_schema = arrow::schema({arrow::field("col1", arrow::int32()),
arrow::field("col2", arrow::int32())});
auto merging_reader =
ColumnMergingReader::Make(target_schema, std::move(readers));
std::shared_ptr<arrow::RecordBatch> result_batch;
arrow::Status status = merging_reader->ReadNext(&result_batch);
EXPECT_FALSE(status.ok());
EXPECT_EQ(status.code(), arrow::StatusCode::Invalid);
}
// Test read error
TEST_F(ColumnMergingReaderTest, ReadError) {
std::vector<std::shared_ptr<arrow::ipc::RecordBatchReader>> readers;
readers.push_back(
std::make_shared<MockErrorRecordBatchReader>(arrow::StatusCode::IOError));
auto merging_reader =
ColumnMergingReader::Make(target_schema_, std::move(readers));
std::shared_ptr<arrow::RecordBatch> result_batch;
arrow::Status status = merging_reader->ReadNext(&result_batch);
EXPECT_FALSE(status.ok());
EXPECT_EQ(status.code(), arrow::StatusCode::IOError);
}
// Test multiple reads
TEST_F(ColumnMergingReaderTest, MultipleReads) {
// Create batches
auto array1a = MakeInt32Array({1, 2}).ValueOrDie();
auto batch1a = arrow::RecordBatch::Make(
arrow::schema({arrow::field("col1", arrow::int32())}), 2, {array1a});
auto array1b = MakeInt32Array({3, 4}).ValueOrDie();
auto batch1b = arrow::RecordBatch::Make(
arrow::schema({arrow::field("col1", arrow::int32())}), 2, {array1b});
auto array2a = MakeInt32Array({5, 6}).ValueOrDie();
auto batch2a = arrow::RecordBatch::Make(
arrow::schema({arrow::field("col2", arrow::int32())}), 2, {array2a});
auto array2b = MakeInt32Array({7, 8}).ValueOrDie();
auto batch2b = arrow::RecordBatch::Make(
arrow::schema({arrow::field("col2", arrow::int32())}), 2, {array2b});
class MultiBatchRecordBatchReader : public arrow::ipc::RecordBatchReader {
public:
explicit MultiBatchRecordBatchReader(
std::vector<std::shared_ptr<arrow::RecordBatch>> batches)
: batches_(std::move(batches)), index_(0) {}
std::shared_ptr<arrow::Schema> schema() const override {
return batches_.empty() ? arrow::schema({}) : batches_[0]->schema();
}
arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch> *out) override {
if (index_ < batches_.size()) {
*out = batches_[index_++];
} else {
*out = nullptr;
}
return arrow::Status::OK();
}
private:
std::vector<std::shared_ptr<arrow::RecordBatch>> batches_;
size_t index_;
};
std::vector<std::shared_ptr<arrow::ipc::RecordBatchReader>> readers;
readers.push_back(std::make_shared<MultiBatchRecordBatchReader>(
std::vector<std::shared_ptr<arrow::RecordBatch>>{batch1a, batch1b}));
readers.push_back(std::make_shared<MultiBatchRecordBatchReader>(
std::vector<std::shared_ptr<arrow::RecordBatch>>{batch2a, batch2b}));
auto target_schema = arrow::schema({arrow::field("col1", arrow::int32()),
arrow::field("col2", arrow::int32())});
auto merging_reader =
ColumnMergingReader::Make(target_schema, std::move(readers));
// First read
std::shared_ptr<arrow::RecordBatch> result_batch;
ASSERT_OK(merging_reader->ReadNext(&result_batch));
ASSERT_NE(result_batch, nullptr);
EXPECT_EQ(result_batch->num_rows(), 2);
auto col1 =
std::static_pointer_cast<arrow::Int32Array>(result_batch->column(0));
auto col2 =
std::static_pointer_cast<arrow::Int32Array>(result_batch->column(1));
EXPECT_EQ(col1->Value(0), 1);
EXPECT_EQ(col1->Value(1), 2);
EXPECT_EQ(col2->Value(0), 5);
EXPECT_EQ(col2->Value(1), 6);
// Second read
ASSERT_OK(merging_reader->ReadNext(&result_batch));
ASSERT_NE(result_batch, nullptr);
EXPECT_EQ(result_batch->num_rows(), 2);
col1 = std::static_pointer_cast<arrow::Int32Array>(result_batch->column(0));
col2 = std::static_pointer_cast<arrow::Int32Array>(result_batch->column(1));
EXPECT_EQ(col1->Value(0), 3);
EXPECT_EQ(col1->Value(1), 4);
EXPECT_EQ(col2->Value(0), 7);
EXPECT_EQ(col2->Value(1), 8);
// Third read - should be EOF
ASSERT_OK(merging_reader->ReadNext(&result_batch));
EXPECT_EQ(result_batch, nullptr);
}
// Test zero row batches
TEST_F(ColumnMergingReaderTest, ZeroRowBatches) {
auto array1 = MakeInt32Array({}).ValueOrDie();
auto batch1 = arrow::RecordBatch::Make(
arrow::schema({arrow::field("col1", arrow::int32())}), 0, {array1});
auto array2 = MakeInt32Array({}).ValueOrDie();
auto batch2 = arrow::RecordBatch::Make(
arrow::schema({arrow::field("col2", arrow::int32())}), 0, {array2});
class MockRecordBatchReader : public arrow::ipc::RecordBatchReader {
public:
explicit MockRecordBatchReader(std::shared_ptr<arrow::RecordBatch> batch)
: batch_(batch), returned_(false) {}
std::shared_ptr<arrow::Schema> schema() const override {
return batch_->schema();
}
arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch> *out) override {
if (!returned_) {
*out = batch_;
returned_ = true;
} else {
*out = nullptr;
}
return arrow::Status::OK();
}
private:
std::shared_ptr<arrow::RecordBatch> batch_;
bool returned_;
};
std::vector<std::shared_ptr<arrow::ipc::RecordBatchReader>> readers;
readers.push_back(std::make_shared<MockRecordBatchReader>(batch1));
readers.push_back(std::make_shared<MockRecordBatchReader>(batch2));
auto target_schema = arrow::schema({arrow::field("col1", arrow::int32()),
arrow::field("col2", arrow::int32())});
auto merging_reader =
ColumnMergingReader::Make(target_schema, std::move(readers));
std::shared_ptr<arrow::RecordBatch> result_batch;
ASSERT_OK(merging_reader->ReadNext(&result_batch));
EXPECT_EQ(result_batch, nullptr);
}
@@ -0,0 +1,805 @@
// 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 "db/index/segment/segment_helper.h"
#include <algorithm>
#include <cstdint>
#include <filesystem>
#include <iostream>
#include <memory>
#include <set>
#include <thread>
#include <variant>
#include <arrow/array/array_binary.h>
#include <arrow/io/file.h>
#include <arrow/ipc/reader.h>
#include <arrow/pretty_print.h>
#include <arrow/result.h>
#include <arrow/table.h>
#include <gtest/gtest.h>
#include "db/common/constants.h"
#include "db/common/file_helper.h"
#include "db/index/column/vector_column/vector_column_indexer.h"
#include "db/index/column/vector_column/vector_column_params.h"
#include "db/index/column/vector_column/vector_index_results.h"
#include "db/index/common/delete_store.h"
#include "db/index/common/id_map.h"
#include "db/index/common/meta.h"
#include "db/index/common/version_manager.h"
#include "db/index/segment/segment.h"
#include "utils/utils.h"
#include "zvec/db/options.h"
#include "zvec/db/query_params.h"
#include "zvec/db/schema.h"
using namespace zvec;
class SegmentHelperTest : public testing::Test {
protected:
void SetUp() override {
ailego::LoggerBroker::SetLevel(ailego::Logger::LEVEL_INFO);
FileHelper::RemoveDirectory(col_path);
FileHelper::CreateDirectory(col_path);
std::string idmap_path =
FileHelper::MakeFilePath(col_path, FileID::ID_FILE, 0);
id_map = IDMap::CreateAndOpen(col_name, idmap_path, true, false);
if (id_map == nullptr) {
throw std::runtime_error("Failed to create id map");
}
std::string delete_store_path =
FileHelper::MakeFilePath(col_path, FileID::DELETE_FILE, 0);
delete_store = std::make_shared<DeleteStore>(col_name);
}
void TearDown() override {
id_map.reset();
delete_store.reset();
// FileHelper::RemoveDirectory(col_path);
}
public:
std::string GetColPath() {
return col_path;
}
protected:
VersionManager::Ptr CreateVersionManager(const CollectionSchema &schema) {
Version version;
version.set_schema(schema);
auto vm = VersionManager::Create(col_path, version);
if (!vm.has_value()) {
throw std::runtime_error("Failed to create version manager");
}
return vm.value();
}
SegmentOptions WriteOptions() const {
return SegmentOptions{false, true, DEFAULT_MAX_BUFFER_SIZE};
}
struct CompactResult {
CompactTask compact_task;
Segment::Ptr output_segment; // null when filter dropped every doc
};
// Execute a CompactTask end-to-end: build it, run it, move the tmp segment
// dir into place, and reopen the output segment in read-only mode.
CompactResult RunCompactAndOpen(CollectionSchema::Ptr schema,
std::vector<Segment::Ptr> segments,
SegmentID output_segment_id,
IndexFilter::Ptr filter,
const VersionManager::Ptr &version_manager,
int concurrency = 1) {
const bool forward_use_parquet = false;
CompactTask task(col_path, schema, std::move(segments), output_segment_id,
std::move(filter), forward_use_parquet, concurrency);
auto segment_task = SegmentTask::CreateCompactTask(task);
EXPECT_NE(segment_task, nullptr);
if (segment_task == nullptr) return {task, nullptr};
auto status = SegmentHelper::Execute(segment_task);
EXPECT_TRUE(status.ok()) << status.message();
auto executed = std::get<CompactTask>(segment_task->task_info());
if (executed.output_segment_meta_ == nullptr) {
return {executed, nullptr};
}
auto tmp_path =
FileHelper::MakeTempSegmentPath(col_path, output_segment_id);
auto dst_path = FileHelper::MakeSegmentPath(col_path, output_segment_id);
EXPECT_TRUE(FileHelper::MoveDirectory(tmp_path, dst_path));
SegmentOptions read_options{true, !forward_use_parquet,
DEFAULT_MAX_BUFFER_SIZE};
version_manager->set_enable_mmap(!forward_use_parquet);
auto seg_ret =
Segment::Open(col_path, *schema, *executed.output_segment_meta_, id_map,
delete_store, version_manager, read_options);
EXPECT_TRUE(seg_ret.has_value());
if (!seg_ret.has_value()) return {executed, nullptr};
return {executed, std::move(seg_ret.value())};
}
std::string col_name = "test_segment_helper";
std::string col_path = "./test_collection";
IDMap::Ptr id_map;
DeleteStore::Ptr delete_store;
};
TEST_F(SegmentHelperTest, CompactTask_General) {
auto schema = test::TestHelper::CreateNormalSchema(false, col_name);
auto version_manager = CreateVersionManager(*schema);
auto write_options = WriteOptions();
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
write_options, 0, 1000);
ASSERT_TRUE(seg1 != nullptr);
ASSERT_TRUE(seg1->flush().ok());
auto seg2 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 1, 1000, id_map, delete_store, version_manager,
write_options, 1000, 1000);
ASSERT_TRUE(seg2 != nullptr);
ASSERT_TRUE(seg2->flush().ok());
SegmentID output_segment_id = 2;
auto [compact_task, seg3] = RunCompactAndOpen(
schema, {seg1, seg2}, output_segment_id, nullptr, version_manager);
ASSERT_NE(seg3, nullptr);
ASSERT_EQ(compact_task.output_segment_meta_->id(), output_segment_id);
ASSERT_FALSE(
compact_task.output_segment_meta_->writing_forward_block().has_value());
ASSERT_EQ(seg3->id(), output_segment_id);
ASSERT_EQ(seg3->doc_count(), seg1->doc_count() + seg2->doc_count());
for (uint64_t i = 0; i < seg3->doc_count(); i++) {
auto doc = seg3->Fetch(i);
ASSERT_NE(doc, nullptr);
auto expect_doc = test::TestHelper::CreateDoc(i, *schema);
ASSERT_EQ(*doc, expect_doc);
}
ASSERT_TRUE(seg1->destroy().ok());
ASSERT_TRUE(seg2->destroy().ok());
}
TEST_F(SegmentHelperTest, CompactTask_ScalarIndex) {
auto schema = test::TestHelper::CreateSchemaWithScalarIndex(false);
auto version_manager = CreateVersionManager(*schema);
auto write_options = WriteOptions();
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
write_options, 0, 1000);
ASSERT_TRUE(seg1 != nullptr);
ASSERT_TRUE(seg1->flush().ok());
auto seg2 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 1, 1000, id_map, delete_store, version_manager,
write_options, 1000, 1000);
ASSERT_TRUE(seg2 != nullptr);
ASSERT_TRUE(seg2->flush().ok());
SegmentID output_segment_id = 2;
auto [compact_task, seg3] = RunCompactAndOpen(
schema, {seg1, seg2}, output_segment_id, nullptr, version_manager);
ASSERT_NE(seg3, nullptr);
ASSERT_EQ(compact_task.output_segment_meta_->id(), output_segment_id);
ASSERT_FALSE(
compact_task.output_segment_meta_->writing_forward_block().has_value());
ASSERT_EQ(seg3->id(), output_segment_id);
ASSERT_EQ(seg3->doc_count(), seg1->doc_count() + seg2->doc_count());
for (uint64_t i = 0; i < seg3->doc_count(); i++) {
auto doc = seg3->Fetch(i);
ASSERT_NE(doc, nullptr);
auto expect_doc = test::TestHelper::CreateDoc(i, *schema);
ASSERT_EQ(*doc, expect_doc);
}
ASSERT_TRUE(seg1->destroy().ok());
ASSERT_TRUE(seg2->destroy().ok());
}
TEST_F(SegmentHelperTest, CompactTask_VectorIndex) {
auto schema = test::TestHelper::CreateSchemaWithVectorIndex();
auto version_manager = CreateVersionManager(*schema);
auto write_options = WriteOptions();
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
write_options, 0, 1000);
ASSERT_TRUE(seg1 != nullptr);
ASSERT_TRUE(seg1->flush().ok());
auto seg2 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 1, 1000, id_map, delete_store, version_manager,
write_options, 1000, 1000);
ASSERT_TRUE(seg2 != nullptr);
ASSERT_TRUE(seg2->flush().ok());
SegmentID output_segment_id = 2;
auto [compact_task, seg3] = RunCompactAndOpen(
schema, {seg1, seg2}, output_segment_id, nullptr, version_manager);
ASSERT_NE(seg3, nullptr);
ASSERT_EQ(compact_task.output_segment_meta_->id(), output_segment_id);
ASSERT_FALSE(
compact_task.output_segment_meta_->writing_forward_block().has_value());
ASSERT_EQ(seg3->id(), output_segment_id);
ASSERT_EQ(seg3->doc_count(), seg1->doc_count() + seg2->doc_count());
for (uint64_t i = 0; i < seg3->doc_count(); i++) {
auto doc = seg3->Fetch(i);
ASSERT_NE(doc, nullptr);
auto expect_doc = test::TestHelper::CreateDoc(i, *schema);
ASSERT_EQ(*doc, expect_doc);
}
ASSERT_TRUE(seg1->destroy().ok());
ASSERT_TRUE(seg2->destroy().ok());
}
TEST_F(SegmentHelperTest, CompactTask_MultipleSegments) {
auto schema = test::TestHelper::CreateNormalSchema(false, col_name);
auto version_manager = CreateVersionManager(*schema);
auto write_options = WriteOptions();
std::vector<Segment::Ptr> input_segs;
const int seg_count = 10;
const int doc_count_per_seg = 100;
for (int i = 0; i < seg_count; i++) {
auto seg = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, i, i * doc_count_per_seg, id_map, delete_store,
version_manager, write_options, i * doc_count_per_seg,
doc_count_per_seg);
ASSERT_TRUE(seg != nullptr);
ASSERT_TRUE(seg->flush().ok());
input_segs.push_back(seg);
}
SegmentID output_segment_id = seg_count;
auto [compact_task, seg3] = RunCompactAndOpen(
schema, input_segs, output_segment_id, nullptr, version_manager);
ASSERT_NE(seg3, nullptr);
ASSERT_EQ(compact_task.output_segment_meta_->id(), output_segment_id);
ASSERT_FALSE(
compact_task.output_segment_meta_->writing_forward_block().has_value());
ASSERT_EQ(seg3->id(), output_segment_id);
ASSERT_EQ(seg3->doc_count(), seg_count * doc_count_per_seg);
for (uint64_t i = 0; i < seg3->doc_count(); i++) {
auto doc = seg3->Fetch(i);
ASSERT_NE(doc, nullptr);
auto expect_doc = test::TestHelper::CreateDoc(i, *schema);
ASSERT_EQ(*doc, expect_doc);
}
}
TEST_F(SegmentHelperTest, CompactTask_Filter) {
auto schema = test::TestHelper::CreateNormalSchema(false, col_name);
auto version_manager = CreateVersionManager(*schema);
auto write_options = WriteOptions();
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
write_options, 0, 1000);
ASSERT_TRUE(seg1 != nullptr);
ASSERT_TRUE(seg1->flush().ok());
auto filter = std::make_shared<EasyIndexFilter>(
[](uint64_t id) -> bool { return id < 10; });
SegmentID output_segment_id = 1;
auto [compact_task, seg2] = RunCompactAndOpen(
schema, {seg1}, output_segment_id, filter, version_manager);
ASSERT_NE(seg2, nullptr);
ASSERT_EQ(compact_task.output_segment_meta_->id(), output_segment_id);
ASSERT_FALSE(
compact_task.output_segment_meta_->writing_forward_block().has_value());
ASSERT_EQ(seg2->id(), output_segment_id);
ASSERT_EQ(seg2->doc_count(), seg1->doc_count() - 10);
ASSERT_TRUE(seg1->destroy().ok());
}
TEST_F(SegmentHelperTest, CompactTask_FilterAll) {
auto schema = test::TestHelper::CreateNormalSchema(false, col_name);
auto version_manager = CreateVersionManager(*schema);
auto write_options = WriteOptions();
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
write_options, 0, 1000);
ASSERT_TRUE(seg1 != nullptr);
ASSERT_TRUE(seg1->flush().ok());
auto filter = std::make_shared<EasyIndexFilter>(
[](uint64_t /*id*/) -> bool { return true; });
SegmentID output_segment_id = 1;
auto [compact_task, output_segment] = RunCompactAndOpen(
schema, {seg1}, output_segment_id, filter, version_manager);
ASSERT_EQ(compact_task.output_segment_meta_, nullptr);
ASSERT_EQ(output_segment, nullptr);
ASSERT_FALSE(FileHelper::DirectoryExists(
FileHelper::MakeTempSegmentPath(col_path, output_segment_id)));
}
TEST_F(SegmentHelperTest, CreateVectorIndexTask_AllFields) {
auto schema = test::TestHelper::CreateNormalSchema(false, col_name);
Version version;
version.set_schema(*schema);
auto version_manager_tmp = VersionManager::Create(col_path, version);
if (!version_manager_tmp.has_value()) {
throw std::runtime_error("Failed to create version manager");
}
auto version_manager = version_manager_tmp.value();
// Create a segment
auto segment = test::TestHelper::CreateSegmentWithDoc(
GetColPath(), *schema, 0, 0, id_map, delete_store, version_manager,
SegmentOptions{false, true, DEFAULT_MAX_BUFFER_SIZE}, 0, 1000);
ASSERT_TRUE(segment != nullptr);
ASSERT_TRUE(segment->dump().ok());
// Create index params
auto index_params =
std::make_shared<HnswIndexParams>(MetricType::L2, // metric_type
16, // m
100 // ef_construction
);
// Create create index task
CreateVectorIndexTask task(
segment,
"", // column_to_build_vector_index (empty means all vector columns)
index_params,
1 // concurrency
);
// Create segment task
auto segment_task = SegmentTask::CreateCreateVectorIndexTask(task);
// Verify task creation
ASSERT_TRUE(segment_task != nullptr);
// Execute the task
Status status = SegmentHelper::Execute(segment_task);
std::cout << "status: " << status.message() << std::endl;
EXPECT_TRUE(status.ok());
// Verify output segment meta
auto index_task = std::get<CreateVectorIndexTask>(segment_task->task_info());
auto output_segment_meta = index_task.output_segment_meta_;
std::cout << "output_segment_meta: "
<< output_segment_meta->to_string_formatted() << std::endl;
ASSERT_EQ(output_segment_meta->id(), 0);
ASSERT_FALSE(output_segment_meta->writing_forward_block().has_value());
auto segment_meta = std::make_shared<SegmentMeta>(*segment->meta());
segment_meta->remove_writing_forward_block();
// create all vector index will not change segment meta
ASSERT_EQ(*output_segment_meta, *segment_meta);
}
TEST_F(SegmentHelperTest, CreateVectorIndexTask_SingleField) {
auto schema = test::TestHelper::CreateNormalSchema(false, col_name);
Version version;
version.set_schema(*schema);
auto version_manager_tmp = VersionManager::Create(col_path, version);
if (!version_manager_tmp.has_value()) {
throw std::runtime_error("Failed to create version manager");
}
auto version_manager = version_manager_tmp.value();
// Create a segment
auto segment = test::TestHelper::CreateSegmentWithDoc(
GetColPath(), *schema, 0, 0, id_map, delete_store, version_manager,
SegmentOptions{false, true, DEFAULT_MAX_BUFFER_SIZE}, 0, 1000);
ASSERT_TRUE(segment != nullptr);
ASSERT_TRUE(segment->dump().ok());
// Create index params
auto index_params =
std::make_shared<HnswIndexParams>(MetricType::IP, // metric_type
16, // m
100 // ef_construction
);
// Create create index task
CreateVectorIndexTask task(segment,
"dense_fp32", // column_to_build_vector_index
// (empty means all vector columns)
index_params,
1 // concurrency
);
// Create segment task
auto segment_task = SegmentTask::CreateCreateVectorIndexTask(task);
// Verify task creation
ASSERT_TRUE(segment_task != nullptr);
// Execute the task
Status status = SegmentHelper::Execute(segment_task);
std::cout << "status: " << status.message() << std::endl;
EXPECT_TRUE(status.ok());
// Verify output segment meta
auto index_task = std::get<CreateVectorIndexTask>(segment_task->task_info());
auto output_segment_meta = index_task.output_segment_meta_;
std::cout << "output_segment_meta: "
<< output_segment_meta->to_string_formatted() << std::endl;
ASSERT_EQ(output_segment_meta->id(), 0);
ASSERT_FALSE(output_segment_meta->writing_forward_block().has_value());
}
TEST_F(SegmentHelperTest, CompactTask_VectorIndexThreeSegmentsRegression) {
auto schema = test::TestHelper::CreateSchemaWithVectorIndex();
auto version_manager = CreateVersionManager(*schema);
auto write_options = WriteOptions();
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
write_options, 0, 300);
auto seg2 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 1, 300, id_map, delete_store, version_manager,
write_options, 300, 300);
auto seg3 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 2, 600, id_map, delete_store, version_manager,
write_options, 600, 300);
ASSERT_TRUE(seg1 != nullptr);
ASSERT_TRUE(seg2 != nullptr);
ASSERT_TRUE(seg3 != nullptr);
ASSERT_TRUE(seg1->flush().ok());
ASSERT_TRUE(seg2->flush().ok());
ASSERT_TRUE(seg3->flush().ok());
auto [compact_task, output_segment] = RunCompactAndOpen(
schema, {seg1, seg2, seg3}, 3, nullptr, version_manager);
ASSERT_NE(output_segment, nullptr);
ASSERT_EQ(output_segment->doc_count(), 900);
ASSERT_NE(output_segment->Fetch(0), nullptr);
ASSERT_NE(output_segment->Fetch(899), nullptr);
}
TEST_F(SegmentHelperTest,
CompactTask_QuantizedVectorIndexThreeSegmentsRegression) {
auto schema = test::TestHelper::CreateSchemaWithVectorIndex(
false, col_name,
std::make_shared<HnswIndexParams>(MetricType::IP, 16, 20,
QuantizeType::FP16));
auto version_manager = CreateVersionManager(*schema);
auto write_options = WriteOptions();
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
write_options, 0, 300);
auto seg2 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 1, 300, id_map, delete_store, version_manager,
write_options, 300, 300);
auto seg3 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 2, 600, id_map, delete_store, version_manager,
write_options, 600, 300);
ASSERT_TRUE(seg1 != nullptr);
ASSERT_TRUE(seg2 != nullptr);
ASSERT_TRUE(seg3 != nullptr);
ASSERT_TRUE(seg1->flush().ok());
ASSERT_TRUE(seg2->flush().ok());
ASSERT_TRUE(seg3->flush().ok());
ASSERT_GT(seg1->get_quant_vector_indexer("dense_fp32").size(), 0u);
ASSERT_GT(seg2->get_quant_vector_indexer("dense_fp32").size(), 0u);
ASSERT_GT(seg3->get_quant_vector_indexer("dense_fp32").size(), 0u);
auto [compact_task, output_segment] = RunCompactAndOpen(
schema, {seg1, seg2, seg3}, 3, nullptr, version_manager);
ASSERT_NE(output_segment, nullptr);
ASSERT_EQ(output_segment->doc_count(), 900);
ASSERT_NE(output_segment->Fetch(0), nullptr);
ASSERT_NE(output_segment->Fetch(899), nullptr);
ASSERT_GT(output_segment->get_vector_indexer("dense_fp32").size(), 0u);
ASSERT_GT(output_segment->get_quant_vector_indexer("dense_fp32").size(), 0u);
}
struct SegmentCompactReuseParam {
IndexParams::Ptr vector_index_params;
IndexType expected_output_type;
};
class SegmentCompactReuseTest
: public SegmentHelperTest,
public testing::WithParamInterface<SegmentCompactReuseParam> {
protected:
// Returns the indexer's underlying VectorIndexParams::type(), defaulting to
// FLAT if the params can't be downcast (matches the freshly-inserted state).
static IndexType IndexerType(const VectorColumnIndexer::Ptr &indexer) {
auto params = std::dynamic_pointer_cast<VectorIndexParams>(
indexer->field_schema().index_params());
return params ? params->type() : IndexType::FLAT;
}
static QuantizeType QuantizeTypeOf(const IndexParams::Ptr &params) {
auto vp = std::dynamic_pointer_cast<VectorIndexParams>(params);
return vp ? vp->quantize_type() : QuantizeType::UNDEFINED;
}
// Run CreateVectorIndexTask on `segment` for `column` with `index_params`,
// then reload the segment so its in-memory indexer reflects the new index
// (matching collection.cc's post-optimize reload path).
void OptimizeSegmentToVectorIndex(const Segment::Ptr &segment,
const CollectionSchema &schema,
const std::string &column,
const IndexParams::Ptr &index_params) {
CreateVectorIndexTask task(segment, column, index_params, 1);
auto segment_task = SegmentTask::CreateCreateVectorIndexTask(task);
ASSERT_NE(segment_task, nullptr);
ASSERT_TRUE(SegmentHelper::Execute(segment_task).ok());
auto executed = std::get<CreateVectorIndexTask>(segment_task->task_info());
ASSERT_NE(executed.output_segment_meta_, nullptr);
ASSERT_TRUE(
segment
->reload_vector_index(schema, executed.output_segment_meta_,
executed.output_vector_indexers_,
executed.output_quant_vector_indexers_)
.ok());
}
struct ScoredDoc {
uint64_t doc_id;
float score;
};
// CreateDoc seeds VECTOR_FP32 with a constant vector of value (doc_id+0.1f).
static std::vector<float> MakeFp32QueryVector(uint64_t doc_id_value,
uint32_t dim) {
return std::vector<float>(dim, static_cast<float>(doc_id_value) + 0.1f);
}
static std::vector<ScoredDoc> RunSearch(
const VectorColumnIndexer::Ptr &indexer, const std::vector<float> &qvec,
uint32_t topk, const zvec::QueryParams::Ptr &query_params) {
vector_column_params::QueryParams qp;
qp.topk = topk;
qp.filter = nullptr;
qp.fetch_vector = false;
qp.query_params = query_params;
vector_column_params::VectorData data{
vector_column_params::DenseVector{qvec.data()}};
auto results = indexer->Search(data, qp);
EXPECT_TRUE(results.has_value());
if (!results.has_value()) return {};
auto vec_res = dynamic_cast<VectorIndexResults *>(results.value().get());
EXPECT_NE(vec_res, nullptr);
if (vec_res == nullptr) return {};
std::vector<ScoredDoc> out;
for (auto it = vec_res->create_iterator(); it->valid(); it->next()) {
out.push_back({it->doc_id(), it->score()});
}
return out;
}
// All test instantiations use IP — higher score is better.
static std::set<uint64_t> MergeTopKIds(
std::vector<std::vector<ScoredDoc>> per_seg, uint32_t topk) {
std::vector<ScoredDoc> all;
for (auto &v : per_seg)
for (auto &d : v) all.push_back(d);
std::sort(all.begin(), all.end(),
[](const ScoredDoc &a, const ScoredDoc &b) {
return a.score > b.score;
});
std::set<uint64_t> ids;
for (size_t i = 0; i < all.size() && ids.size() < topk; ++i) {
ids.insert(all[i].doc_id);
}
return ids;
}
static zvec::QueryParams::Ptr MakeIsLinearQueryParam(IndexType type) {
switch (type) {
case IndexType::HNSW: {
auto p = std::make_shared<zvec::HnswQueryParams>();
p->set_is_linear(true);
return p;
}
case IndexType::IVF: {
auto p = std::make_shared<zvec::IVFQueryParams>();
p->set_is_linear(true);
return p;
}
case IndexType::HNSW_RABITQ: {
auto p = std::make_shared<zvec::HnswRabitqQueryParams>();
p->set_is_linear(true);
return p;
}
case IndexType::FLAT:
default:
return std::make_shared<zvec::FlatQueryParams>();
}
}
};
// Mimic the normal insertion lifecycle: small segments accumulate vectors
// in flat storage (no vector index built yet), then compaction merges them
// into a single segment whose vector column is built per schema.
TEST_P(SegmentCompactReuseTest, OptimizedSegmentsReuseFirstIndexer) {
const auto &param = GetParam();
auto schema = test::TestHelper::CreateSchemaWithVectorIndex(
false, col_name, param.vector_index_params);
auto version_manager = CreateVersionManager(*schema);
auto write_options = WriteOptions();
constexpr int kSegCount = 3;
constexpr int kDocsPerSeg = 300;
constexpr uint32_t kTopK = 10;
constexpr uint32_t kDim = 128;
std::vector<Segment::Ptr> segs;
for (int i = 0; i < kSegCount; i++) {
auto seg = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, i, i * kDocsPerSeg, id_map, delete_store,
version_manager, write_options, i * kDocsPerSeg, kDocsPerSeg);
ASSERT_NE(seg, nullptr);
ASSERT_TRUE(seg->flush().ok());
segs.push_back(seg);
}
// Capture groundtruth via FlatQuery on each source segment while every
// segment is still backed by a flat indexer (before seg[0] is optimized).
const std::vector<uint64_t> query_doc_values{0, kDocsPerSeg,
kSegCount * kDocsPerSeg - 1};
std::vector<std::set<uint64_t>> groundtruth;
groundtruth.reserve(query_doc_values.size());
auto flat_qp = std::make_shared<zvec::FlatQueryParams>();
for (uint64_t qv : query_doc_values) {
auto qvec = MakeFp32QueryVector(qv, kDim);
std::vector<std::vector<ScoredDoc>> per_seg;
per_seg.reserve(segs.size());
// Per-segment indexers use block-local doc ids (0..kDocsPerSeg-1).
// The compacted output indexer reindexes them sequentially across
// segments, so add segment offset to align id spaces before merging.
for (size_t s = 0; s < segs.size(); ++s) {
auto in_indexers = segs[s]->get_vector_indexer("dense_fp32");
ASSERT_FALSE(in_indexers.empty());
ASSERT_EQ(IndexerType(in_indexers.front()), IndexType::FLAT);
auto local = RunSearch(in_indexers.front(), qvec, kTopK, flat_qp);
const uint64_t offset = static_cast<uint64_t>(s) * kDocsPerSeg;
for (auto &d : local) d.doc_id += offset;
per_seg.push_back(std::move(local));
}
auto gt = MergeTopKIds(std::move(per_seg), kTopK);
ASSERT_EQ(gt.size(), kTopK);
groundtruth.push_back(std::move(gt));
}
// Optimize seg[0]'s vector fields to the parametric index type, mimicking
// the lifecycle the compact path exercises.
for (const auto &vf : schema->vector_fields()) {
OptimizeSegmentToVectorIndex(segs[0], *schema, vf->name(),
vf->index_params());
}
// For quantized index types (e.g. HNSW_RABITQ) the built index lives in
// get_quant_vector_indexer; get_vector_indexer keeps the raw FLAT
// indexer. See CompactTask_QuantizedVectorIndexThreeSegmentsRegression.
const bool quantized =
QuantizeTypeOf(param.vector_index_params) != QuantizeType::UNDEFINED;
for (int i = 0; i < kSegCount; i++) {
auto in_indexers = quantized
? segs[i]->get_quant_vector_indexer("dense_fp32")
: segs[i]->get_vector_indexer("dense_fp32");
ASSERT_FALSE(in_indexers.empty());
ASSERT_EQ(IndexerType(in_indexers.front()),
i == 0 ? param.expected_output_type : IndexType::FLAT);
}
auto [compact_task, output_segment] =
RunCompactAndOpen(schema, segs, kSegCount, nullptr, version_manager);
ASSERT_NE(output_segment, nullptr);
ASSERT_EQ(output_segment->doc_count(), kSegCount * kDocsPerSeg);
ASSERT_NE(output_segment->Fetch(0), nullptr);
ASSERT_NE(output_segment->Fetch(kSegCount * kDocsPerSeg - 1), nullptr);
auto out_indexers =
quantized ? output_segment->get_quant_vector_indexer("dense_fp32")
: output_segment->get_vector_indexer("dense_fp32");
ASSERT_FALSE(out_indexers.empty());
EXPECT_EQ(IndexerType(out_indexers.front()), param.expected_output_type);
// is_linear queries on the merged indexer must reproduce the pre-compact
// groundtruth. Quantized indexers are allowed a small recall hit.
auto linear_qp = MakeIsLinearQueryParam(param.expected_output_type);
const double kMinRecall = quantized ? 0.8 : 1.0;
for (size_t qi = 0; qi < query_doc_values.size(); ++qi) {
auto qvec = MakeFp32QueryVector(query_doc_values[qi], kDim);
auto hits = RunSearch(out_indexers.front(), qvec, kTopK, linear_qp);
ASSERT_EQ(hits.size(), kTopK);
size_t intersect = 0;
for (const auto &h : hits) {
if (groundtruth[qi].count(h.doc_id)) intersect++;
}
double recall = static_cast<double>(intersect) / kTopK;
EXPECT_GE(recall, kMinRecall)
<< "query[" << qi << "] (value=" << query_doc_values[qi]
<< ") recall=" << recall;
}
}
INSTANTIATE_TEST_SUITE_P(Hnsw, SegmentCompactReuseTest,
testing::Values(SegmentCompactReuseParam{
std::make_shared<HnswIndexParams>(MetricType::IP,
16, 200),
IndexType::HNSW}));
// CreateNormalSchema() only puts the test's vector_index_params on dense_fp32.
// The other 4 vector fields are hardcoded — dense_fp16/dense_int8/sparse_fp16
// are always FlatIndexParams, and sparse_fp32 gets the
// cloned params only if supports_sparse is true (utils.cc:117-124), which
// excludes IVF and HNSW_RABITQ — so for IVF it also falls back to FLAT.
INSTANTIATE_TEST_SUITE_P(
Ivf, SegmentCompactReuseTest,
testing::Values(SegmentCompactReuseParam{
std::make_shared<IVFIndexParams>(MetricType::IP, 10, 4, false,
QuantizeType::UNDEFINED),
IndexType::IVF}));
#if RABITQ_SUPPORTED
INSTANTIATE_TEST_SUITE_P(HnswRabitq, SegmentCompactReuseTest,
testing::Values(SegmentCompactReuseParam{
std::make_shared<HnswRabitqIndexParams>(
MetricType::IP, 7, 256, 16, 200, 0),
IndexType::HNSW_RABITQ}));
#endif
TEST_F(SegmentHelperTest, CompactTask_FilterMultiSegmentsRegression) {
auto schema = test::TestHelper::CreateSchemaWithVectorIndex();
auto version_manager = CreateVersionManager(*schema);
auto write_options = WriteOptions();
auto seg1 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 0, 0, id_map, delete_store, version_manager,
write_options, 0, 400);
auto seg2 = test::TestHelper::CreateSegmentWithDoc(
col_path, *schema, 1, 400, id_map, delete_store, version_manager,
write_options, 400, 400);
ASSERT_TRUE(seg1 != nullptr);
ASSERT_TRUE(seg2 != nullptr);
ASSERT_TRUE(seg1->flush().ok());
ASSERT_TRUE(seg2->flush().ok());
auto filter = std::make_shared<EasyIndexFilter>(
[](uint64_t id) -> bool { return id < 100 || (id >= 400 && id < 450); });
auto [compact_task, output_segment] =
RunCompactAndOpen(schema, {seg1, seg2}, 2, filter, version_manager);
ASSERT_NE(output_segment, nullptr);
ASSERT_EQ(output_segment->doc_count(), 650);
}
@@ -0,0 +1,234 @@
// 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 <cstdint>
#include <memory>
#include <string>
#include <vector>
#include <arrow/array.h>
#include <arrow/record_batch.h>
#include <gtest/gtest.h>
#include "db/common/constants.h"
#include "segment_test_fixture.h"
using namespace zvec;
namespace {
struct LocalRowIdProjection {
std::vector<std::string> columns;
int row_id_column_index;
};
const std::vector<LocalRowIdProjection> &LocalRowIdProjections() {
static const std::vector<LocalRowIdProjection> projections = {
{{LOCAL_ROW_ID, "id", "name"}, 0},
{{"id", LOCAL_ROW_ID, "name"}, 1},
{{"id", "name", LOCAL_ROW_ID}, 2},
};
return projections;
}
void ExpectSingleRowLocalRowID(const ExecBatchPtr &batch, int column_index,
uint64_t expected) {
ASSERT_TRUE(batch != nullptr);
EXPECT_EQ(batch->length, 1);
EXPECT_EQ(batch->values.size(), 3);
auto local_row_id_scalar = batch->values[column_index].scalar();
ASSERT_TRUE(local_row_id_scalar != nullptr);
auto local_row_id_value =
std::dynamic_pointer_cast<arrow::UInt64Scalar>(local_row_id_scalar);
ASSERT_TRUE(local_row_id_value != nullptr);
EXPECT_EQ(local_row_id_value->value, expected);
}
void ExpectScanLocalRowIDColumn(const RecordBatchReaderPtr &reader,
int column_index, uint32_t expected_rows) {
ASSERT_TRUE(reader != nullptr);
ASSERT_TRUE(reader->schema() != nullptr);
std::shared_ptr<arrow::RecordBatch> batch;
uint32_t total_doc = 0;
while (true) {
auto status = reader->ReadNext(&batch);
ASSERT_TRUE(status.ok()) << status.ToString();
if (batch == nullptr) break;
ASSERT_GT(batch->num_columns(), column_index);
EXPECT_EQ(batch->column(column_index)->type()->id(), arrow::Type::UINT64);
EXPECT_EQ(batch->column_name(column_index), LOCAL_ROW_ID);
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, expected_rows);
}
void ExpectFetchedLocalRowIDs(
const TablePtr &table, int column_index,
const std::vector<int> &expected_segment_doc_ids) {
ASSERT_TRUE(table != nullptr);
EXPECT_EQ(table->num_columns(), 3);
EXPECT_EQ(table->num_rows(),
static_cast<int64_t>(expected_segment_doc_ids.size()));
auto field = table->schema()->field(column_index);
EXPECT_EQ(field->name(), LOCAL_ROW_ID);
auto id_column = table->column(column_index);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(id_column->chunk(0));
ASSERT_TRUE(id_array != nullptr);
std::vector<uint64_t> actual_ids;
actual_ids.reserve(id_array->length());
for (int64_t i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
std::vector<uint64_t> expected_u64_ids(expected_segment_doc_ids.begin(),
expected_segment_doc_ids.end());
EXPECT_EQ(actual_ids, expected_u64_ids)
<< "LOCAL_ROW_ID values don't match expected order";
}
} // namespace
TEST_P(SegmentTest, FetchSingleRowWithLocalRowIDInRequestedPosition) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_,
options_, 0, 10);
ASSERT_TRUE(segment != nullptr);
for (const auto &projection : LocalRowIdProjections()) {
ExecBatchPtr batch = segment->fetch(projection.columns, 4);
ExpectSingleRowLocalRowID(batch, projection.row_id_column_index, 4);
}
}
TEST_P(SegmentTest, ScanAndFetchPreserveLocalRowIDColumnPosition) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_,
options_, 0, 10);
ASSERT_TRUE(segment != nullptr);
std::vector<int> segment_doc_ids = {0, 3, 6, 1, 0};
for (const auto &projection : LocalRowIdProjections()) {
auto reader = segment->scan(projection.columns);
ExpectScanLocalRowIDColumn(reader, projection.row_id_column_index, 10);
auto table = segment->fetch(projection.columns, segment_doc_ids);
ExpectFetchedLocalRowIDs(table, projection.row_id_column_index,
segment_doc_ids);
}
}
TEST_P(SegmentTest, ScanLocalRowIDIsSegmentLocal) {
options_.max_buffer_size_ = 1 * 1024;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 100, id_map_, delete_store_, version_manager_,
options_, 100, 25);
ASSERT_TRUE(segment != nullptr);
auto reader = segment->scan({LOCAL_ROW_ID, "id"});
ASSERT_TRUE(reader != nullptr);
std::shared_ptr<arrow::RecordBatch> batch;
std::vector<uint64_t> actual_ids;
while (true) {
auto status = reader->ReadNext(&batch);
ASSERT_TRUE(status.ok()) << status.ToString();
if (batch == nullptr) break;
ASSERT_EQ(batch->num_columns(), 2);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(batch->column(0));
ASSERT_TRUE(id_array != nullptr);
for (int64_t i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
}
std::vector<uint64_t> expected_ids;
for (uint64_t i = 0; i < 25; ++i) {
expected_ids.push_back(i);
}
EXPECT_EQ(actual_ids, expected_ids);
}
TEST_P(SegmentTest, ScanOnlyLocalRowIDDoesNotExposeGlobalDocID) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 100, id_map_, delete_store_, version_manager_,
options_, 100, 10);
ASSERT_TRUE(segment != nullptr);
auto reader = segment->scan({LOCAL_ROW_ID});
ASSERT_TRUE(reader != nullptr);
ASSERT_TRUE(reader->schema() != nullptr);
ASSERT_EQ(reader->schema()->num_fields(), 1);
EXPECT_EQ(reader->schema()->field(0)->name(), LOCAL_ROW_ID);
std::shared_ptr<arrow::RecordBatch> batch;
std::vector<uint64_t> actual_ids;
while (true) {
auto status = reader->ReadNext(&batch);
ASSERT_TRUE(status.ok()) << status.ToString();
if (batch == nullptr) break;
ASSERT_EQ(batch->num_columns(), 1);
EXPECT_EQ(batch->column_name(0), LOCAL_ROW_ID);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(batch->column(0));
ASSERT_TRUE(id_array != nullptr);
for (int64_t i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
}
std::vector<uint64_t> expected_ids;
for (uint64_t i = 0; i < 10; ++i) {
expected_ids.push_back(i);
}
EXPECT_EQ(actual_ids, expected_ids);
}
TEST_P(SegmentTest, DocCountDeleteFilterWithNonZeroGlobalDocID) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 100, id_map_, delete_store_, version_manager_,
options_, 0, 10);
ASSERT_TRUE(segment != nullptr);
auto status = segment->Delete("pk_5");
EXPECT_TRUE(status.ok()) << "Delete by pk failed: " << status.message();
status = segment->Delete(103);
EXPECT_TRUE(status.ok()) << "Delete by global doc id failed: "
<< status.message();
EXPECT_EQ(segment->doc_count(), 10);
EXPECT_EQ(segment->doc_count(delete_store_->make_filter()), 8);
}
INSTANTIATE_TEST_SUITE_P(MMapTest, SegmentTest, testing::Values(true, false));
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,105 @@
// 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.
#pragma once
#include <memory>
#include <stdexcept>
#include <string>
#include <gtest/gtest.h>
#include <zvec/ailego/buffer/block_eviction_queue.h>
#include "db/common/file_helper.h"
#include "db/index/common/delete_store.h"
#include "db/index/common/id_map.h"
#include "db/index/common/version_manager.h"
#include "utils/utils.h"
#include "zvec/db/options.h"
namespace zvec {
class SegmentTest : public testing::TestWithParam<bool> {
protected:
void SetUp() override {
ailego::LoggerBroker::SetLevel(ailego::Logger::LEVEL_WARN);
zvec::ailego::MemoryLimitPool::get_instance().init(MIN_MEMORY_LIMIT_BYTES);
FileHelper::RemoveDirectory(col_path_);
FileHelper::CreateDirectory(col_path_);
auto id_map_path = FileHelper::MakeFilePath(col_path_, FileID::ID_FILE, 0);
id_map_ = IDMap::CreateAndOpen(col_name_, id_map_path, true, false);
if (id_map_ == nullptr) {
throw std::runtime_error("Failed to create id map");
}
delete_store_ = std::make_shared<DeleteStore>(col_name_);
schema_ =
test::TestHelper::CreateSchemaWithScalarIndex(false, false, col_name_);
schema_->add_field(
std::make_shared<FieldSchema>("id", DataType::INT32, false));
schema_->add_field(
std::make_shared<FieldSchema>("name", DataType::STRING, false));
schema_->add_field(
std::make_shared<FieldSchema>("age", DataType::UINT32, false));
schema_->add_field(
std::make_shared<FieldSchema>("binary", DataType::BINARY, false));
schema_->add_field(std::make_shared<FieldSchema>(
"array_binary", DataType::ARRAY_BINARY, false));
bool enable_mmap = GetParam();
Version version;
version.set_schema(*schema_);
version.set_enable_mmap(enable_mmap);
auto version_manager_tmp = VersionManager::Create(col_path_, version);
if (!version_manager_tmp.has_value()) {
throw std::runtime_error("Failed to create version manager");
}
version_manager_ = version_manager_tmp.value();
options_.read_only_ = false;
options_.enable_mmap_ = enable_mmap;
options_.max_buffer_size_ = 64 * 1024 * 1024;
}
void TearDown() override {
id_map_.reset();
delete_store_.reset();
version_manager_.reset();
FileHelper::RemoveDirectory(col_path_);
}
public:
std::string GetColPath() {
return col_path_;
}
protected:
std::string col_name_ = "test_segment";
std::string col_path_ = "./test_collection";
IDMap::Ptr id_map_;
DeleteStore::Ptr delete_store_;
VersionManager::Ptr version_manager_;
CollectionSchema::Ptr schema_;
SegmentOptions options_;
};
} // namespace zvec
@@ -0,0 +1,290 @@
// 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 "db/index/segment/sql_expr_parser.h"
#include <arrow/array.h>
#include <arrow/compute/api.h>
#include <arrow/dataset/api.h>
#include <arrow/dataset/discovery.h>
#include <arrow/memory_pool.h>
#include <arrow/result.h>
#include <arrow/table.h>
#include <arrow/testing/gtest_util.h>
#include <gtest/gtest.h>
#include "utils/utils.h"
using namespace arrow;
using namespace arrow::dataset;
using namespace zvec;
class SqlExprParserTest : public ::testing::Test {
protected:
void SetUp() override {
// Setup code if needed
}
void TearDown() override {
// Cleanup code if needed
}
};
TEST_F(SqlExprParserTest, ParseAllSupportedTypes) {
auto schema = arrow::schema({arrow::field("int32", arrow::int32()),
arrow::field("uint32", arrow::uint32()),
arrow::field("float", arrow::float32()),
arrow::field("double", arrow::float64()),
arrow::field("int64", arrow::int64()),
arrow::field("uint64", arrow::uint64()),
arrow::field("string", arrow::utf8()),
arrow::field("bool", arrow::boolean())});
EXPECT_TRUE(ParseToExpression("int32 + uint32", schema).ok());
EXPECT_TRUE(ParseToExpression("float * double", schema).ok());
EXPECT_TRUE(ParseToExpression("int64 - uint64", schema).ok());
EXPECT_TRUE(ParseToExpression("int32 / float", schema).ok());
EXPECT_TRUE(ParseToExpression("double + int64", schema).ok());
EXPECT_TRUE(ParseToExpression("uint32 * int32", schema).ok());
EXPECT_TRUE(ParseToExpression("int32 + float - double", schema).ok());
EXPECT_TRUE(ParseToExpression("int64 * uint32 / float", schema).ok());
EXPECT_TRUE(ParseToExpression("(int32 + float) * double", schema).ok());
EXPECT_TRUE(
ParseToExpression("int32 + (float - double) * int64", schema).ok());
EXPECT_TRUE(
ParseToExpression("((int32 + uint32) * float) - (double / int64)", schema)
.ok());
EXPECT_TRUE(ParseToExpression("int32 + 100", schema).ok());
EXPECT_TRUE(ParseToExpression("float * 3.14", schema).ok());
EXPECT_TRUE(ParseToExpression("double - 2.5", schema).ok());
EXPECT_TRUE(ParseToExpression("(int64 + 10) * (uint32 - 5)", schema).ok());
EXPECT_TRUE(ParseToExpression("-int32", schema).ok());
EXPECT_TRUE(ParseToExpression("-(float + double)", schema).ok());
}
TEST_F(SqlExprParserTest, ParseStringExpression) {
auto schema = arrow::schema({arrow::field("name", arrow::utf8()),
arrow::field("age", arrow::int32())});
auto result = ParseToExpression("name = 'John'", schema);
EXPECT_FALSE(result.ok());
}
TEST_F(SqlExprParserTest, ParseBooleanExpression) {
auto schema = arrow::schema({arrow::field("active", arrow::boolean()),
arrow::field("age", arrow::int32())});
auto result = ParseToExpression("active AND age > 18", schema);
EXPECT_FALSE(result.ok());
}
TEST_F(SqlExprParserTest, ParseListExpression) {
auto schema = arrow::schema(
{arrow::field("int32_list", arrow::list(arrow::int32())),
arrow::field("float64_list", arrow::list(arrow::float64())),
arrow::field("int32", arrow::int32()),
arrow::field("float64", arrow::float64())});
auto result = ParseToExpression("int32 + int32_list", schema);
EXPECT_FALSE(result.ok());
result = ParseToExpression("float64 + float64_list", schema);
EXPECT_FALSE(result.ok());
}
TEST_F(SqlExprParserTest, ParseComplexExpression) {
auto schema = arrow::schema({arrow::field("price", arrow::float64()),
arrow::field("quantity", arrow::int32()),
arrow::field("discount", arrow::float64())});
auto result = ParseToExpression("price * quantity * (1 - discount)", schema);
EXPECT_TRUE(result.ok()) << "Failed to parse SQL expression status:"
<< result.status().ToString();
}
TEST_F(SqlExprParserTest, ParseInvalidExpression) {
auto schema = arrow::schema({arrow::field("a", arrow::int32())});
auto result = ParseToExpression("a + ", schema);
EXPECT_FALSE(result.ok());
}
TEST_F(SqlExprParserTest, ParseNonExistentField) {
auto schema = arrow::schema({arrow::field("a", arrow::int32())});
auto result = ParseToExpression("b + 1", schema);
EXPECT_FALSE(result.ok());
}
TEST_F(SqlExprParserTest, ParseFunctionCall) {
auto schema = arrow::schema({arrow::field("value", arrow::float64())});
auto result = ParseToExpression("sqrt(value)", schema);
EXPECT_FALSE(result.ok());
}
TEST_F(SqlExprParserTest, ParseComplexCombinations) {
auto schema = arrow::schema(
{arrow::field("a", arrow::int32()), arrow::field("b", arrow::float64()),
arrow::field("c", arrow::int64()), arrow::field("d", arrow::float32())});
// Deeply nested expressions
auto result = ParseToExpression("((a + b) * (c - d)) / (a + 1)", schema);
EXPECT_TRUE(result.ok()) << "Failed to parse SQL expression status:"
<< result.status().ToString();
// Multi-level parentheses expressions
result = ParseToExpression("(((a + b) - c) * d) + (a / b)", schema);
EXPECT_TRUE(result.ok()) << "Failed to parse SQL expression status:"
<< result.status().ToString();
// Mixed constants and variables
result = ParseToExpression("(a + 10) * (b - 2.5) / (c + 100)", schema);
EXPECT_TRUE(result.ok()) << "Failed to parse SQL expression status:"
<< result.status().ToString();
}
// Test negative number expressions
TEST_F(SqlExprParserTest, ParseNegativeNumbers) {
auto schema = arrow::schema({arrow::field("id", arrow::int32()),
arrow::field("value", arrow::float64())});
// Test negative fields
auto result = ParseToExpression("-id", schema);
EXPECT_TRUE(result.ok()) << "Failed to parse SQL expression status:"
<< result.status().ToString();
// Test negative numbers combined with other operators
result = ParseToExpression("-id + value", schema);
EXPECT_TRUE(result.ok()) << "Failed to parse SQL expression status:"
<< result.status().ToString();
// Test nested negative expressions
result = ParseToExpression("-(-id)", schema);
EXPECT_TRUE(result.ok()) << "Failed to parse SQL expression status:"
<< result.status().ToString();
// Test complex negative expressions
result = ParseToExpression("-(id + value) * 2", schema);
EXPECT_TRUE(result.ok()) << "Failed to parse SQL expression status:"
<< result.status().ToString();
}
// Create a simple Table
std::shared_ptr<arrow::Table> MakeTestTable() {
// Create integer column
arrow::Int32Builder int_builder;
ARROW_EXPECT_OK(int_builder.AppendValues({1, 2, 3, 4, 5}));
std::shared_ptr<arrow::Array> int_array;
ARROW_EXPECT_OK(int_builder.Finish(&int_array));
// Create double column
arrow::DoubleBuilder double_builder;
ARROW_EXPECT_OK(double_builder.AppendValues({1.1, 2.2, 3.3, 4.4, 5.5}));
std::shared_ptr<arrow::Array> double_array;
ARROW_EXPECT_OK(double_builder.Finish(&double_array));
// Create string column
arrow::StringBuilder string_builder;
ARROW_EXPECT_OK(string_builder.Append("a"));
ARROW_EXPECT_OK(string_builder.Append("b"));
ARROW_EXPECT_OK(string_builder.Append("c"));
ARROW_EXPECT_OK(string_builder.Append("d"));
ARROW_EXPECT_OK(string_builder.Append("e"));
std::shared_ptr<arrow::Array> string_array;
ARROW_EXPECT_OK(string_builder.Finish(&string_array));
// Create boolean column
arrow::BooleanBuilder bool_builder;
ARROW_EXPECT_OK(bool_builder.Append(true));
ARROW_EXPECT_OK(bool_builder.Append(false));
ARROW_EXPECT_OK(bool_builder.Append(true));
ARROW_EXPECT_OK(bool_builder.Append(false));
ARROW_EXPECT_OK(bool_builder.Append(true));
std::shared_ptr<arrow::Array> bool_array;
ARROW_EXPECT_OK(bool_builder.Finish(&bool_array));
// Build table
auto schema = arrow::schema({arrow::field("int_col", arrow::int32()),
arrow::field("double_col", arrow::float64()),
arrow::field("string_col", arrow::utf8()),
arrow::field("bool_col", arrow::boolean())});
auto int_chunked = std::make_shared<arrow::ChunkedArray>(int_array);
auto double_chunked = std::make_shared<arrow::ChunkedArray>(double_array);
auto string_chunked = std::make_shared<arrow::ChunkedArray>(string_array);
auto bool_chunked = std::make_shared<arrow::ChunkedArray>(bool_array);
return arrow::Table::Make(
schema, {int_chunked, double_chunked, string_chunked, bool_chunked});
}
// Convert Table to Dataset (for testing)
arrow::Result<std::shared_ptr<arrow::dataset::Dataset>> MakeTestDataset(
const std::shared_ptr<arrow::Table> &table) {
return std::make_shared<arrow::dataset::InMemoryDataset>(table);
}
TEST_F(SqlExprParserTest, ParseAndScanDataSet) {
auto status = arrow::compute::Initialize();
auto schema = arrow::schema({arrow::field("int_col", arrow::int32()),
arrow::field("double_col", arrow::float64()),
arrow::field("string_col", arrow::utf8()),
arrow::field("bool_col", arrow::boolean())});
// Step 1: Create test table
auto table = MakeTestTable();
// Step 2: Convert to Dataset
auto dataset = MakeTestDataset(table).ValueOrDie();
// Step 3: Create scanner and project expression A + B
auto scanner_builder = dataset->NewScan().ValueOrDie();
auto expr = ParseToExpression("int_col + double_col", schema).ValueOrDie();
status = scanner_builder->Project({expr}, {"sum"});
auto scanner = scanner_builder->Finish().ValueOrDie();
// Step 4: Execute and get results
auto result_table = scanner->ToTable().ValueOrDie();
ASSERT_TRUE(result_table != nullptr);
ASSERT_EQ(result_table->num_rows(), 5);
auto int_col = table->column(0); // int_col
auto double_col = table->column(1); // double_col
auto sum_col = result_table->column(0); // sum column
for (int64_t i = 0; i < table->num_rows(); ++i) {
auto int_value =
std::static_pointer_cast<arrow::Int32Array>(int_col->chunk(0))
->Value(i);
auto double_value =
std::static_pointer_cast<arrow::DoubleArray>(double_col->chunk(0))
->Value(i);
auto sum_value =
std::static_pointer_cast<arrow::DoubleArray>(sum_col->chunk(0))
->Value(i);
ASSERT_NEAR(int_value + double_value, sum_value, 1e-10);
}
}
@@ -0,0 +1,145 @@
// 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 <arrow/array.h>
#include <arrow/builder.h>
#include <arrow/dataset/api.h>
#include <arrow/table.h>
#include <arrow/type.h>
#include <gmock/gmock-matchers.h>
#include <gtest/gtest.h>
#include "db/index/segment/sql_expr_parser.h"
using arrow::Status;
using arrow::compute::Expression;
namespace compute = arrow::compute;
using namespace zvec;
arrow::Result<Expression> ParseAndValidate(
const std::string &expr, const std::shared_ptr<arrow::Schema> &schema) {
ARROW_ASSIGN_OR_RAISE(auto parsed, ParseToExpression(expr, schema));
return CheckSupportedArithmeticExpression(parsed, *schema);
}
class ExprValidatorTest : public ::testing::Test {
protected:
void SetUp() override {
schema_ = arrow::schema({arrow::field("int32_col", arrow::int32()),
arrow::field("double_col", arrow::float64()),
arrow::field("str_col", arrow::utf8())});
std::vector<std::shared_ptr<arrow::Array>> arrays;
for (const auto &field : schema_->fields()) {
std::unique_ptr<arrow::ArrayBuilder> builder;
ASSERT_TRUE(arrow::MakeBuilder(arrow::default_memory_pool(),
field->type(), &builder)
.ok());
std::shared_ptr<arrow::Array> array;
ASSERT_TRUE(builder->Finish(&array).ok());
arrays.push_back(array);
}
auto table = arrow::Table::Make(schema_, arrays);
dataset_ = std::make_shared<arrow::dataset::InMemoryDataset>(table);
}
std::shared_ptr<arrow::Schema> schema_;
std::shared_ptr<arrow::dataset::Dataset> dataset_;
};
TEST_F(ExprValidatorTest, SingleNumericColumn_Valid) {
auto result = ParseAndValidate("int32_col", schema_);
EXPECT_TRUE(result.ok()) << result.status().ToString();
result = ParseAndValidate("double_col", schema_);
EXPECT_TRUE(result.ok()) << result.status().ToString();
}
TEST_F(ExprValidatorTest, UnaryPositive_Supported) {
auto result = ParseAndValidate("+int32_col", schema_);
EXPECT_TRUE(result.ok()) << result.status().ToString();
result = ParseAndValidate("+double_col", schema_);
EXPECT_TRUE(result.ok()) << result.status().ToString();
}
TEST_F(ExprValidatorTest, UnaryNegative_Supported) {
auto result = ParseAndValidate("-int32_col", schema_);
EXPECT_TRUE(result.ok()) << result.status().ToString();
result = ParseAndValidate("-double_col", schema_);
EXPECT_TRUE(result.ok()) << result.status().ToString();
}
TEST_F(ExprValidatorTest, Binary_Op_With_Literal_Valid) {
auto result = ParseAndValidate("int32_col + 1", schema_);
EXPECT_TRUE(result.ok()) << result.status().ToString();
result = ParseAndValidate("int32_col - 100", schema_);
EXPECT_TRUE(result.ok());
result = ParseAndValidate("1.5 * double_col", schema_);
EXPECT_TRUE(result.ok());
result = ParseAndValidate("double_col / 2.0", schema_);
EXPECT_TRUE(result.ok());
result = ParseAndValidate("100 - int32_col", schema_);
EXPECT_TRUE(result.ok());
}
TEST_F(ExprValidatorTest, NonNumericColumn_Rejected) {
auto result = ParseAndValidate("str_col", schema_);
EXPECT_FALSE(result.ok());
EXPECT_THAT(result.status().ToString(), ::testing::HasSubstr("not numeric"));
result = ParseAndValidate("+str_col", schema_);
EXPECT_FALSE(result.ok());
EXPECT_THAT(result.status().ToString(), ::testing::HasSubstr("not numeric"));
result = ParseAndValidate("-str_col", schema_);
EXPECT_FALSE(result.ok());
EXPECT_THAT(result.status().ToString(), ::testing::HasSubstr("not numeric"));
}
TEST_F(ExprValidatorTest, TwoColumns_Operations_Rejected) {
auto result = ParseAndValidate("int32_col + double_col", schema_);
EXPECT_FALSE(result.ok());
result = ParseAndValidate("int32_col + int32_col", schema_);
EXPECT_FALSE(result.ok());
}
TEST_F(ExprValidatorTest, PureLiteral_Rejected) {
auto result = ParseAndValidate("123", schema_);
EXPECT_TRUE(result.ok());
result = ParseAndValidate("+123", schema_);
EXPECT_TRUE(result.ok());
result = ParseAndValidate("-456", schema_);
EXPECT_TRUE(result.ok()) << result.status().ToString();
}
TEST_F(ExprValidatorTest, NestedExpression_Rejected) {
auto result = ParseAndValidate("(int32_col + 1)", schema_);
EXPECT_TRUE(result.ok()) << result.status().ToString();
}
TEST_F(ExprValidatorTest, InvalidFunctionOrSyntax) {
auto result = ParseAndValidate("int32_col || 'abc'", schema_);
EXPECT_FALSE(result.ok());
result = ParseAndValidate("sqrt(int32_col)", schema_);
EXPECT_FALSE(result.ok());
}
@@ -0,0 +1,130 @@
// 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 "db/index/storage/arrow_ipc_writer.h"
#include <iostream>
#include <arrow/array/builder_primitive.h>
#include <arrow/ipc/reader.h>
#include <arrow/record_batch.h>
#include <arrow/status.h>
#include <gtest/gtest.h>
#include "db/index/storage/store_helper.h"
using namespace zvec;
auto schema = arrow::schema(
{arrow::field("id", arrow::int32()), arrow::field("name", arrow::utf8())});
std::shared_ptr<arrow::RecordBatchReader> CreateTestReader(int start_id,
int count) {
arrow::Int32Builder id_builder;
arrow::StringBuilder name_builder;
arrow::Status s;
for (int i = 0; i < count; ++i) {
s = id_builder.Append(start_id + i);
if (!s.ok()) {
return nullptr;
}
s = name_builder.Append("User" + std::to_string(start_id + i));
if (!s.ok()) {
return nullptr;
}
}
std::shared_ptr<arrow::Array> id_array, name_array;
s = id_builder.Finish(&id_array);
if (!s.ok()) {
return nullptr;
}
s = name_builder.Finish(&name_array);
if (!s.ok()) {
return nullptr;
}
auto batch = arrow::RecordBatch::Make(schema, count, {id_array, name_array});
auto maybe_reader = arrow::RecordBatchReader::Make({batch}, schema);
if (!maybe_reader.ok()) {
return nullptr;
}
return *maybe_reader;
}
TEST(ArrowIpcWriter, General) {
std::string output_file_path = "output.ipc";
ArrowIpcWriter writer(output_file_path);
// writer.SetMaxRowsPerGroup(1000); // 可选:控制每组行数
// 第一次插入
{
auto reader1 = CreateTestReader(1, 3);
ASSERT_NE(reader1, nullptr);
auto status = writer.insert(reader1);
ASSERT_TRUE(status.ok());
std::cout << "Inserted batch 1" << std::endl;
}
// 第二次插入
{
auto reader2 = CreateTestReader(4, 2);
ASSERT_NE(reader2, nullptr);
auto status = writer.insert(reader2);
ASSERT_TRUE(status.ok());
std::cout << "Inserted batch 2" << std::endl;
}
// 第三次插入
{
auto reader3 = CreateTestReader(6, 4);
ASSERT_NE(reader3, nullptr);
auto status = writer.insert(reader3);
ASSERT_TRUE(status.ok());
std::cout << "Inserted batch 3" << std::endl;
}
// 最后关闭文件
auto status = writer.finalize();
if (!status.ok()) {
std::cerr << "Finalize failed: " << status.ToString() << std::endl;
}
std::cout << "Parquet file written successfully to output.parquet"
<< std::endl;
// 读取文件
std::shared_ptr<arrow::io::RandomAccessFile> output_file_;
std::string output_file_path_cp;
auto as = CreateRandomAccessFileByUri(output_file_path, &output_file_,
&output_file_path_cp);
ASSERT_TRUE(as.ok());
auto result = arrow::ipc::RecordBatchFileReader::Open(output_file_);
ASSERT_TRUE(result.ok());
auto reader = std::move(result).ValueOrDie();
ASSERT_EQ(reader->num_record_batches(), 3);
int num_rows = 0;
for (int i = 0; i < reader->num_record_batches(); i++) {
std::shared_ptr<arrow::RecordBatch> batch;
auto res = reader->ReadRecordBatch(i);
ASSERT_TRUE(res.ok());
batch = std::move(res).ValueOrDie();
num_rows += batch->num_rows();
}
ASSERT_EQ(num_rows, 9);
}
@@ -0,0 +1,441 @@
// 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 <cstdint>
#include <cstdlib>
#include <filesystem>
#include <iostream>
#include <memory>
#include <thread>
#include <arrow/api.h>
#include <arrow/result.h>
#include <arrow/table.h>
#include <gtest/gtest.h>
#include "db/index/storage/bufferpool_forward_store.h"
#include "utils/utils.h"
using namespace zvec;
class BufferPoolStoreTest : public testing::Test {
protected:
void SetUp() override {
auto s = test::TestHelper::WriteTestFile(parquet_path, FileFormat::PARQUET);
if (!s.ok()) {
std::cout << "err: " << s.message() << std::endl;
exit(1);
}
zvec::ailego::MemoryLimitPool::get_instance().init(10 * 1024 * 1024);
}
void TearDown() override {
if (std::filesystem::exists(parquet_path)) {
std::filesystem::remove(parquet_path);
}
}
std::string parquet_path = "test.parquet";
};
TEST_F(BufferPoolStoreTest, ParquetFetch) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
TablePtr table = store->fetch({"id", "name", "score"}, {0, 1, 2});
ASSERT_TRUE(table != nullptr);
EXPECT_EQ(table->num_rows(), 3);
}
TEST_F(BufferPoolStoreTest, ParquetFetchWithSelectColumns) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
TablePtr table = store->fetch({"id", "name"}, {0, 1, 2});
ASSERT_TRUE(table != nullptr);
EXPECT_EQ(table->num_rows(), 3);
EXPECT_EQ(table->num_columns(), 2);
}
TEST_F(BufferPoolStoreTest, ParquetFetchWithUID) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
auto table = store->fetch({USER_ID, "id", "name"}, {0, 1, 2});
ASSERT_TRUE(table != nullptr);
EXPECT_EQ(table->num_rows(), 3);
EXPECT_EQ(table->num_columns(), 3);
}
TEST_F(BufferPoolStoreTest, ParquetFetchWithGlobalDocID) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
auto table = store->fetch({GLOBAL_DOC_ID, "id", "name"}, {0, 1, 2});
ASSERT_TRUE(table != nullptr);
EXPECT_EQ(table->num_rows(), 3);
EXPECT_EQ(table->num_columns(), 3);
}
TEST_F(BufferPoolStoreTest, ParquetFetchWitEmptyColumns) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
TablePtr table = store->fetch({}, std::vector<int>{});
EXPECT_EQ(table, nullptr);
}
TEST_F(BufferPoolStoreTest, ParquetFetchWitEmptyIndices) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
TablePtr table = store->fetch({"id", "name"}, std::vector<int>{});
ASSERT_TRUE(table != nullptr);
EXPECT_EQ(table->num_rows(), 0);
EXPECT_EQ(table->num_columns(), 2);
}
TEST_F(BufferPoolStoreTest, ParquetFetchWithMoreIndices) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
TablePtr table = store->fetch({"id"}, {0, 1, 2, 3, 6, 2, 1, 7});
ASSERT_TRUE(table != nullptr);
EXPECT_EQ(table->num_rows(), 8);
EXPECT_EQ(table->num_columns(), 1);
}
TEST_F(BufferPoolStoreTest, ParquetFetchWithInvalidIndices) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
TablePtr table = store->fetch({"id"}, {0, 1, 30});
ASSERT_TRUE(table == nullptr);
}
TEST_F(BufferPoolStoreTest, ParquetFetchCheckOrderWithLocalRowIDMiddle) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
TablePtr table =
store->fetch({"id", "name", LOCAL_ROW_ID, "score"}, {0, 3, 6, 1, 0});
ASSERT_TRUE(table != nullptr);
EXPECT_EQ(table->num_rows(), 5);
EXPECT_EQ(table->num_columns(), 4);
auto field = table->schema()->field(2);
EXPECT_EQ(field->name(), LOCAL_ROW_ID);
// Get data from the _zvec_row_id_ column for each row
auto id_column = table->column(2);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(id_column->chunk(0));
std::vector<int32_t> expected_ids = {0, 3, 6, 1, 0};
std::vector<int32_t> actual_ids;
for (int i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
EXPECT_EQ(actual_ids, expected_ids)
<< "ID column values don't match expected order";
}
TEST_F(BufferPoolStoreTest, ParquetFetchCheckOrderWithLocalRowIDEnd) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
TablePtr table =
store->fetch({"id", "name", "score", LOCAL_ROW_ID}, {0, 3, 6, 1, 0});
ASSERT_TRUE(table != nullptr);
EXPECT_EQ(table->num_rows(), 5);
EXPECT_EQ(table->num_columns(), 4);
auto field = table->schema()->field(3);
EXPECT_EQ(field->name(), LOCAL_ROW_ID);
// Get data from the _zvec_row_id_ column for each row
auto id_column = table->column(3);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(id_column->chunk(0));
std::vector<int32_t> expected_ids = {0, 3, 6, 1, 0};
std::vector<int32_t> actual_ids;
for (int i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
EXPECT_EQ(actual_ids, expected_ids)
<< "ID column values don't match expected order";
}
TEST_F(BufferPoolStoreTest, ParquetScan) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
auto reader = store->scan({"id", "name", "score"});
int batch_count = 0;
int total_rows = 0;
while (true) {
std::shared_ptr<arrow::RecordBatch> batch;
auto status = reader->ReadNext(&batch);
ASSERT_TRUE(status.ok());
if (batch == nullptr) {
break;
}
EXPECT_GT(batch->num_rows(), 0);
EXPECT_EQ(batch->num_columns(), 3);
batch_count++;
total_rows += batch->num_rows();
}
EXPECT_GT(batch_count, 0);
EXPECT_EQ(total_rows, 10);
}
TEST_F(BufferPoolStoreTest, ParquetScanWithSelectColumns) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
auto reader = store->scan({"id", "name"});
int batch_count = 0;
int total_rows = 0;
while (true) {
std::shared_ptr<arrow::RecordBatch> batch;
auto status = reader->ReadNext(&batch);
ASSERT_TRUE(status.ok());
if (batch == nullptr) {
break;
}
EXPECT_GT(batch->num_rows(), 0);
EXPECT_EQ(batch->num_columns(), 2);
batch_count++;
total_rows += batch->num_rows();
}
EXPECT_GT(batch_count, 0);
EXPECT_EQ(total_rows, 10);
}
TEST_F(BufferPoolStoreTest, ParquetScanWithInvalidColumn) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
auto reader = store->scan({"id", "unknown_column"});
ASSERT_TRUE(reader == nullptr);
}
TEST_F(BufferPoolStoreTest, ParquetScanWithUserID) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
auto reader = store->scan({USER_ID, "id", "name", "score"});
int batch_count = 0;
int total_rows = 0;
while (true) {
std::shared_ptr<arrow::RecordBatch> batch;
auto status = reader->ReadNext(&batch);
ASSERT_TRUE(status.ok());
if (batch == nullptr) {
break;
}
EXPECT_GT(batch->num_rows(), 0);
EXPECT_EQ(batch->num_columns(), 4);
batch_count++;
total_rows += batch->num_rows();
}
EXPECT_GT(batch_count, 0);
EXPECT_EQ(total_rows, 10);
}
TEST_F(BufferPoolStoreTest, ParquetScanWithGlobalDocID) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
auto reader = store->scan({GLOBAL_DOC_ID, "id", "name", "score"});
int batch_count = 0;
int total_rows = 0;
while (true) {
std::shared_ptr<arrow::RecordBatch> batch;
auto status = reader->ReadNext(&batch);
ASSERT_TRUE(status.ok());
if (batch == nullptr) {
break;
}
EXPECT_GT(batch->num_rows(), 0);
EXPECT_EQ(batch->num_columns(), 4);
batch_count++;
total_rows += batch->num_rows();
}
EXPECT_GT(batch_count, 0);
EXPECT_EQ(total_rows, 10);
}
TEST_F(BufferPoolStoreTest, ParquetFetchSingleRow) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
ExecBatchPtr batch = store->fetch({"id", "name", "score"}, 0);
ASSERT_TRUE(batch != nullptr);
EXPECT_EQ(batch->length, 1);
EXPECT_EQ(batch->values.size(), 3);
auto id_scalar = batch->values[0].scalar();
ASSERT_TRUE(id_scalar != nullptr);
auto id_value = std::dynamic_pointer_cast<arrow::Int32Scalar>(id_scalar);
ASSERT_TRUE(id_value != nullptr);
EXPECT_EQ(id_value->value, 1);
}
TEST_F(BufferPoolStoreTest, ParquetFetchSpecificRow) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
ExecBatchPtr batch = store->fetch({"id", "name", "score"}, 3);
ASSERT_TRUE(batch != nullptr);
EXPECT_EQ(batch->length, 1);
EXPECT_EQ(batch->values.size(), 3);
auto id_scalar = batch->values[0].scalar();
ASSERT_TRUE(id_scalar != nullptr);
auto id_value = std::dynamic_pointer_cast<arrow::Int32Scalar>(id_scalar);
ASSERT_TRUE(id_value != nullptr);
EXPECT_EQ(id_value->value, 4);
}
TEST_F(BufferPoolStoreTest, ParquetFetchSingleRowWithUserID) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
ExecBatchPtr batch = store->fetch({USER_ID, "id", "name"}, 1);
ASSERT_TRUE(batch != nullptr);
EXPECT_EQ(batch->length, 1);
EXPECT_EQ(batch->values.size(), 3);
auto user_id_scalar = batch->values[0].scalar();
ASSERT_TRUE(user_id_scalar != nullptr);
EXPECT_TRUE(std::dynamic_pointer_cast<arrow::StringScalar>(user_id_scalar) !=
nullptr);
}
TEST_F(BufferPoolStoreTest, ParquetFetchSingleRowWithGlobalDocID) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
ExecBatchPtr batch = store->fetch({GLOBAL_DOC_ID, "id", "name"}, 4);
ASSERT_TRUE(batch != nullptr);
EXPECT_EQ(batch->length, 1);
EXPECT_EQ(batch->values.size(), 3);
auto global_doc_id_scalar = batch->values[0].scalar();
ASSERT_TRUE(global_doc_id_scalar != nullptr);
EXPECT_TRUE(std::dynamic_pointer_cast<arrow::UInt64Scalar>(
global_doc_id_scalar) != nullptr);
}
TEST_F(BufferPoolStoreTest, ParquetFetchSingleRowWithNegativeIndex) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
ExecBatchPtr batch = store->fetch({"id", "name"}, -1);
EXPECT_EQ(batch, nullptr);
}
TEST_F(BufferPoolStoreTest, ParquetFetchSingleRowWithOutOfRangeIndex) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
ExecBatchPtr batch = store->fetch({"id", "name"}, 15);
EXPECT_EQ(batch, nullptr);
}
TEST_F(BufferPoolStoreTest, ParquetFetchSingleRowWithInvalidColumn) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
ExecBatchPtr batch = store->fetch({"id", "invalid_column"}, 0);
EXPECT_EQ(batch, nullptr);
}
TEST_F(BufferPoolStoreTest, ParquetFetchSingleRowWithEmptyColumns) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
ExecBatchPtr batch = store->fetch({}, 0);
EXPECT_EQ(batch, nullptr);
}
TEST_F(BufferPoolStoreTest, AllDataTypeFetchSingleRow) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
EXPECT_TRUE(store->Open().ok());
ExecBatchPtr batch = store->fetch({"id", "list_int32"}, 2);
ASSERT_TRUE(batch != nullptr);
EXPECT_EQ(batch->length, 1);
EXPECT_EQ(batch->values.size(), 2);
auto id_scalar = batch->values[0].scalar();
ASSERT_TRUE(id_scalar != nullptr);
auto id_value = std::dynamic_pointer_cast<arrow::Int32Scalar>(id_scalar);
ASSERT_TRUE(id_value != nullptr);
EXPECT_EQ(id_value->value, 3);
auto list_scalar = batch->values[1].scalar();
ASSERT_TRUE(list_scalar != nullptr);
auto list_value = std::dynamic_pointer_cast<arrow::ListScalar>(list_scalar);
ASSERT_TRUE(list_value != nullptr);
EXPECT_EQ(list_value->value->length(), 128);
auto list_array =
std::dynamic_pointer_cast<arrow::Int32Array>(list_value->value);
ASSERT_TRUE(list_array != nullptr);
for (int i = 0; i < 10 && i < list_array->length(); ++i) {
EXPECT_EQ(list_array->Value(i), 2 * 10 + i);
}
}
TEST_F(BufferPoolStoreTest, AllDataType) {
auto mmap_store = std::make_shared<BufferPoolForwardStore>(parquet_path);
ASSERT_TRUE(mmap_store->Open().ok());
std::vector<std::string> columns = {"id", "list_int32"};
std::vector<int> indices = {0, 3, 6, 1, 0};
TablePtr mmap_table = mmap_store->fetch(columns, indices);
ASSERT_TRUE(mmap_table != nullptr);
EXPECT_EQ(mmap_table->num_rows(), 5);
EXPECT_EQ(mmap_table->num_columns(), 2);
for (size_t j = 0; j < columns.size(); ++j) {
auto column = mmap_table->column(j);
for (int k = 0; k < column->num_chunks(); ++k) {
auto array = column->chunk(k);
if (array->type()->id() == arrow::Type::INT32) {
auto int_array = std::static_pointer_cast<arrow::Int32Array>(array);
for (int i = 0; i < array->length(); ++i) {
int32_t value = int_array->Value(i);
EXPECT_EQ(value, indices[i] + 1);
}
} else if (array->type()->id() == arrow::Type::LIST) {
auto list_array = std::static_pointer_cast<arrow::ListArray>(array);
for (int i = 0; i < array->length(); ++i) {
auto list_value = list_array->value_slice(i);
auto list_value_array =
std::static_pointer_cast<arrow::Int32Array>(list_value);
EXPECT_EQ(list_value_array->length(), 128);
for (int m = 0; m < list_value_array->length(); ++m) {
int32_t value = list_value_array->Value(m);
EXPECT_EQ(value, indices[i] * 10 + m);
}
}
}
}
}
}
TEST_F(BufferPoolStoreTest, DeleteDestructs) {
BufferPoolForwardStore *store = new BufferPoolForwardStore(parquet_path);
delete store;
}
TEST_F(BufferPoolStoreTest, PhysicSchema) {
auto store = std::make_shared<BufferPoolForwardStore>(parquet_path);
ASSERT_NE(store, nullptr);
EXPECT_TRUE(store->Open().ok());
EXPECT_NE(store->physic_schema(), nullptr);
}
File diff suppressed because it is too large Load Diff
+676
View File
@@ -0,0 +1,676 @@
// 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 <cstdint>
#include <filesystem>
#include <iostream>
#include <memory>
#include <thread>
#include <arrow/api.h>
#include <arrow/result.h>
#include <arrow/table.h>
#include <gtest/gtest.h>
#include "db/common/constants.h"
#define private public
#define protected public
#include "db/index/storage/mmap_forward_store.h"
#undef private
#undef protected
#include "utils/utils.h"
using namespace zvec;
class MmapStoreTest : public testing::Test {
protected:
void SetUp() override {
auto s = test::TestHelper::WriteTestFile(ipc_path, FileFormat::IPC);
if (!s.ok()) {
std::cout << s.message() << std::endl;
exit(1);
}
s = test::TestHelper::WriteTestFile(parquet_path, FileFormat::PARQUET);
if (!s.ok()) {
std::cout << s.message() << std::endl;
exit(1);
}
}
void TearDown() override {
if (std::filesystem::exists(ipc_path)) {
std::filesystem::remove(ipc_path);
}
if (std::filesystem::exists(parquet_path)) {
std::filesystem::remove(parquet_path);
}
}
std::string ipc_path = "test.ipc";
std::string parquet_path = "test.parquet";
};
TEST_F(MmapStoreTest, GeneralIPC) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
TablePtr ipc_table =
ipc_store->fetch({"id", "name", "score"}, {0, 3, 6, 1, 0});
ASSERT_TRUE(ipc_table != nullptr);
EXPECT_EQ(ipc_table->num_rows(), 5);
auto table_reader = ipc_store->scan({"id", "name", "score"});
int batch_count = 0;
while (true) {
std::shared_ptr<arrow::RecordBatch> batch;
auto status = table_reader->ReadNext(&batch);
ASSERT_TRUE(status.ok());
if (batch == nullptr) {
break;
}
ASSERT_GT(batch->num_rows(), 0);
batch_count++;
}
ASSERT_EQ(batch_count, 4);
}
TEST_F(MmapStoreTest, IPCFetchWithLocalRowID) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
TablePtr ipc_table =
ipc_store->fetch({LOCAL_ROW_ID, "id", "name", "score"}, {0, 3, 6, 1, 0});
ASSERT_TRUE(ipc_table != nullptr);
EXPECT_EQ(ipc_table->num_columns(), 4);
EXPECT_EQ(ipc_table->num_rows(), 5);
}
TEST_F(MmapStoreTest, IPCCheckOrderWithLocalRowIDMiddle) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
TablePtr mmap_table =
ipc_store->fetch({"id", "name", LOCAL_ROW_ID, "score"}, {0, 3, 6, 1, 0});
ASSERT_TRUE(mmap_table != nullptr);
EXPECT_EQ(mmap_table->num_rows(), 5);
EXPECT_EQ(mmap_table->num_columns(), 4);
auto field = mmap_table->schema()->field(2);
EXPECT_EQ(field->name(), LOCAL_ROW_ID);
// Get data from the _zvec_row_id_ column for each row
auto id_column = mmap_table->column(2);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(id_column->chunk(0));
std::vector<int32_t> expected_ids = {0, 3, 6, 1, 0};
std::vector<int32_t> actual_ids;
for (int i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
EXPECT_EQ(actual_ids, expected_ids)
<< "ID column values don't match expected order";
}
TEST_F(MmapStoreTest, IPCCheckOrderWithLocalRowIDEnd) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
TablePtr mmap_table =
ipc_store->fetch({"id", "name", "score", LOCAL_ROW_ID}, {0, 3, 6, 1, 0});
ASSERT_TRUE(mmap_table != nullptr);
EXPECT_EQ(mmap_table->num_rows(), 5);
EXPECT_EQ(mmap_table->num_columns(), 4);
auto field = mmap_table->schema()->field(3);
EXPECT_EQ(field->name(), LOCAL_ROW_ID);
// Get data from the _zvec_row_id_ column for each row
auto id_column = mmap_table->column(3);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(id_column->chunk(0));
std::vector<int32_t> expected_ids = {0, 3, 6, 1, 0};
std::vector<int32_t> actual_ids;
for (int i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
EXPECT_EQ(actual_ids, expected_ids)
<< "ID column values don't match expected order";
}
TEST_F(MmapStoreTest, IPCFetchWithUID) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
TablePtr ipc_table =
ipc_store->fetch({USER_ID, "id", "name", "score"}, {0, 3, 6, 1, 0});
ASSERT_TRUE(ipc_table != nullptr);
EXPECT_EQ(ipc_table->num_columns(), 4);
EXPECT_EQ(ipc_table->num_rows(), 5);
}
TEST_F(MmapStoreTest, IPCFetchWithGlobalDocID) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
TablePtr ipc_table =
ipc_store->fetch({GLOBAL_DOC_ID, "id", "name", "score"}, {0, 3, 6, 1, 0});
ASSERT_TRUE(ipc_table != nullptr);
EXPECT_EQ(ipc_table->num_columns(), 4);
EXPECT_EQ(ipc_table->num_rows(), 5);
}
TEST_F(MmapStoreTest, IPCFetchWithEmptyColumns) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
TablePtr ipc_table = ipc_store->fetch({}, std::vector<int>{});
EXPECT_EQ(ipc_table, nullptr);
}
TEST_F(MmapStoreTest, IPCFetchWithInvalidColumns) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
TablePtr ipc_table =
ipc_store->fetch({"id", "unknown_column"}, std::vector<int>{});
EXPECT_EQ(ipc_table, nullptr);
}
TEST_F(MmapStoreTest, IPCFetchWithEmptyIndices) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
TablePtr ipc_table =
ipc_store->fetch({"id", "name", "score"}, std::vector<int>{});
ASSERT_TRUE(ipc_table != nullptr);
EXPECT_EQ(ipc_table->num_rows(), 0);
EXPECT_EQ(ipc_table->num_columns(), 3);
}
TEST_F(MmapStoreTest, IPCFetchWithInvalidIndices) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
TablePtr ipc_table =
ipc_store->fetch({"id"}, std::vector<int>{-1}); // Negative index
EXPECT_EQ(ipc_table, nullptr);
ipc_table =
ipc_store->fetch({"id"}, std::vector<int>{100}); // Out of range index
EXPECT_EQ(ipc_table, nullptr);
}
TEST_F(MmapStoreTest, IPCFetchWithEmptyColumnsValidIndices) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
TablePtr ipc_table = ipc_store->fetch({}, {0, 1});
EXPECT_EQ(ipc_table, nullptr);
}
TEST_F(MmapStoreTest, IPCScan) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
auto table_reader = ipc_store->scan({"id", "name", "score"});
ASSERT_TRUE(table_reader != nullptr);
EXPECT_NE(table_reader->schema(), nullptr);
int batch_count = 0;
int total_rows = 0;
while (true) {
std::shared_ptr<arrow::RecordBatch> batch;
auto status = table_reader->ReadNext(&batch);
ASSERT_TRUE(status.ok());
if (batch == nullptr) {
break;
}
EXPECT_GT(batch->num_rows(), 0);
EXPECT_EQ(batch->num_columns(), 3);
batch_count++;
total_rows += batch->num_rows();
}
EXPECT_GT(batch_count, 0);
EXPECT_EQ(total_rows, 10);
}
TEST_F(MmapStoreTest, IPCScanWithSelectColumns) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
auto table_reader = ipc_store->scan({"id", "name"});
int batch_count = 0;
int total_rows = 0;
while (true) {
std::shared_ptr<arrow::RecordBatch> batch;
auto status = table_reader->ReadNext(&batch);
ASSERT_TRUE(status.ok());
if (batch == nullptr) {
break;
}
EXPECT_GT(batch->num_rows(), 0);
EXPECT_EQ(batch->num_columns(), 2);
batch_count++;
total_rows += batch->num_rows();
}
EXPECT_GT(batch_count, 0);
EXPECT_EQ(total_rows, 10);
}
TEST_F(MmapStoreTest, IPCScanWithInvalidColumn) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
auto table_reader = ipc_store->scan({"id", "unknown_column"});
ASSERT_TRUE(table_reader == nullptr);
}
TEST_F(MmapStoreTest, IPCScanWithUserID) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
auto table_reader = ipc_store->scan({USER_ID, "id", "name", "score"});
int batch_count = 0;
int total_rows = 0;
while (true) {
std::shared_ptr<arrow::RecordBatch> batch;
auto status = table_reader->ReadNext(&batch);
ASSERT_TRUE(status.ok());
if (batch == nullptr) {
break;
}
EXPECT_GT(batch->num_rows(), 0);
EXPECT_EQ(batch->num_columns(), 4);
batch_count++;
total_rows += batch->num_rows();
}
EXPECT_GT(batch_count, 0);
EXPECT_EQ(total_rows, 10);
}
TEST_F(MmapStoreTest, IPCScanWithGlobalDocID) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
auto table_reader = ipc_store->scan({GLOBAL_DOC_ID, "id", "name", "score"});
int batch_count = 0;
int total_rows = 0;
while (true) {
std::shared_ptr<arrow::RecordBatch> batch;
auto status = table_reader->ReadNext(&batch);
ASSERT_TRUE(status.ok());
if (batch == nullptr) {
break;
}
EXPECT_GT(batch->num_rows(), 0);
EXPECT_EQ(batch->num_columns(), 4);
batch_count++;
total_rows += batch->num_rows();
}
EXPECT_GT(batch_count, 0);
EXPECT_EQ(total_rows, 10);
}
TEST_F(MmapStoreTest, GeneralParquet) {
auto mmap_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(mmap_store->Open().ok());
TablePtr mmap_table = mmap_store->fetch({"id", "name", "score"}, {0, 1, 2});
ASSERT_TRUE(mmap_table != nullptr);
EXPECT_EQ(mmap_table->num_rows(), 3);
EXPECT_EQ(mmap_table->num_columns(), 3);
}
TEST_F(MmapStoreTest, ParquetFetchWitEmptyColumns) {
auto mmap_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(mmap_store->Open().ok());
TablePtr mmap_table = mmap_store->fetch({}, std::vector<int>{});
EXPECT_EQ(mmap_table, nullptr);
}
TEST_F(MmapStoreTest, ParquetFetchWithInvalidIndices) {
auto parquet_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(parquet_store->Open().ok());
TablePtr parquet_table =
parquet_store->fetch({"id"}, std::vector<int>{-1}); // Negative index
EXPECT_EQ(parquet_table, nullptr);
parquet_table = parquet_store->fetch(
{"id"}, std::vector<int>{100}); // Out of range index
EXPECT_EQ(parquet_table, nullptr);
}
TEST_F(MmapStoreTest, ParquetCheckOrder) {
auto mmap_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(mmap_store->Open().ok());
TablePtr mmap_table =
mmap_store->fetch({"id", "name", "score"}, {0, 3, 6, 1, 0});
ASSERT_TRUE(mmap_table != nullptr);
EXPECT_EQ(mmap_table->num_rows(), 5);
EXPECT_EQ(mmap_table->num_columns(), 3);
// Get data from the id column for each row
auto id_column = mmap_table->column(0); // id column is the first column
auto id_array =
std::dynamic_pointer_cast<arrow::Int32Array>(id_column->chunk(0));
std::vector<int32_t> expected_ids = {
1, 4, 7, 2, 1}; // Corresponding to indices 0, 3, 6, 1, 0
std::vector<int32_t> actual_ids;
for (int i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
EXPECT_EQ(actual_ids, expected_ids)
<< "ID column values don't match expected order";
}
TEST_F(MmapStoreTest, ParquetCheckOrderWithLocalRowIDMiddle) {
auto mmap_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(mmap_store->Open().ok());
TablePtr mmap_table =
mmap_store->fetch({"id", "name", LOCAL_ROW_ID, "score"}, {0, 3, 6, 1, 0});
ASSERT_TRUE(mmap_table != nullptr);
EXPECT_EQ(mmap_table->num_rows(), 5);
EXPECT_EQ(mmap_table->num_columns(), 4);
auto field = mmap_table->schema()->field(2);
EXPECT_EQ(field->name(), LOCAL_ROW_ID);
// Get data from the _zvec_row_id_ column for each row
auto id_column = mmap_table->column(2);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(id_column->chunk(0));
std::vector<int32_t> expected_ids = {0, 3, 6, 1, 0};
std::vector<int32_t> actual_ids;
for (int i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
EXPECT_EQ(actual_ids, expected_ids)
<< "ID column values don't match expected order";
}
TEST_F(MmapStoreTest, ParquetCheckOrderWithLocalRowIDEnd) {
auto mmap_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(mmap_store->Open().ok());
TablePtr mmap_table =
mmap_store->fetch({"id", "name", "score", LOCAL_ROW_ID}, {0, 3, 6, 1, 0});
ASSERT_TRUE(mmap_table != nullptr);
EXPECT_EQ(mmap_table->num_rows(), 5);
EXPECT_EQ(mmap_table->num_columns(), 4);
auto field = mmap_table->schema()->field(3);
EXPECT_EQ(field->name(), LOCAL_ROW_ID);
// Get data from the _zvec_row_id_ column for each row
auto id_column = mmap_table->column(3);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(id_column->chunk(0));
std::vector<int32_t> expected_ids = {0, 3, 6, 1, 0};
std::vector<int32_t> actual_ids;
for (int i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
EXPECT_EQ(actual_ids, expected_ids)
<< "ID column values don't match expected order";
}
TEST_F(MmapStoreTest, ParquetScan) {
auto mmap_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(mmap_store->Open().ok());
auto table_reader = mmap_store->scan({"id", "name", "score"});
ASSERT_TRUE(table_reader != nullptr);
EXPECT_NE(table_reader->schema(), nullptr);
int batch_count = 0;
int total_rows = 0;
while (true) {
std::shared_ptr<arrow::RecordBatch> batch;
auto status = table_reader->ReadNext(&batch);
ASSERT_TRUE(status.ok());
if (batch == nullptr) {
break;
}
EXPECT_GT(batch->num_rows(), 0);
EXPECT_EQ(batch->num_columns(), 3);
batch_count++;
total_rows += batch->num_rows();
}
EXPECT_GT(batch_count, 0);
EXPECT_EQ(total_rows, 10);
}
TEST_F(MmapStoreTest, ParquetScanWithInvalidColumn) {
auto mmap_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(mmap_store->Open().ok());
auto table_reader = mmap_store->scan({"id", "unknown_column"});
ASSERT_TRUE(table_reader == nullptr);
}
TEST_F(MmapStoreTest, ParquetScanWithUserID) {
auto mmap_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(mmap_store->Open().ok());
auto table_reader = mmap_store->scan({USER_ID, "id", "name", "score"});
int batch_count = 0;
int total_rows = 0;
while (true) {
std::shared_ptr<arrow::RecordBatch> batch;
auto status = table_reader->ReadNext(&batch);
ASSERT_TRUE(status.ok());
if (batch == nullptr) {
break;
}
EXPECT_GT(batch->num_rows(), 0);
EXPECT_EQ(batch->num_columns(), 4);
batch_count++;
total_rows += batch->num_rows();
}
EXPECT_GT(batch_count, 0);
EXPECT_EQ(total_rows, 10);
}
TEST_F(MmapStoreTest, ParquetScanWithGlobalDocID) {
auto mmap_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(mmap_store->Open().ok());
auto table_reader = mmap_store->scan({GLOBAL_DOC_ID, "id", "name", "score"});
int batch_count = 0;
int total_rows = 0;
while (true) {
std::shared_ptr<arrow::RecordBatch> batch;
auto status = table_reader->ReadNext(&batch);
ASSERT_TRUE(status.ok());
if (batch == nullptr) {
break;
}
EXPECT_GT(batch->num_rows(), 0);
EXPECT_EQ(batch->num_columns(), 4);
batch_count++;
total_rows += batch->num_rows();
}
EXPECT_GT(batch_count, 0);
EXPECT_EQ(total_rows, 10);
}
TEST_F(MmapStoreTest, IPCFetchSingleRow) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
auto func = [&](int index) -> void {
ExecBatchPtr ipc_batch = ipc_store->fetch({"id", "name", "score"}, index);
ASSERT_TRUE(ipc_batch != nullptr);
EXPECT_EQ(ipc_batch->length, 1);
EXPECT_EQ(ipc_batch->values.size(), 3);
auto id_scalar = ipc_batch->values[0].scalar();
auto name_scalar = ipc_batch->values[1].scalar();
auto score_scalar = ipc_batch->values[2].scalar();
EXPECT_EQ(std::dynamic_pointer_cast<arrow::Int32Scalar>(id_scalar)->value,
index + 1);
};
for (size_t i = 0; i < 10; i++) {
func(i);
}
}
TEST_F(MmapStoreTest, ParquetFetchSingleRow) {
auto parquet_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(parquet_store->Open().ok());
auto func = [&](int index) -> void {
ExecBatchPtr parquet_batch =
parquet_store->fetch({"id", "name", "score"}, index);
ASSERT_TRUE(parquet_batch != nullptr);
EXPECT_EQ(parquet_batch->length, 1);
EXPECT_EQ(parquet_batch->values.size(), 3);
auto id_scalar = parquet_batch->values[0].scalar();
auto name_scalar = parquet_batch->values[1].scalar();
auto score_scalar = parquet_batch->values[2].scalar();
EXPECT_EQ(std::dynamic_pointer_cast<arrow::Int32Scalar>(id_scalar)->value,
index + 1);
};
for (size_t i = 0; i < 10; i++) {
func(i);
}
}
TEST_F(MmapStoreTest, IPCFetchSingleRowWithInvalidIndex) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
ExecBatchPtr ipc_batch = ipc_store->fetch({"id", "name"}, -1);
EXPECT_EQ(ipc_batch, nullptr);
ipc_batch = ipc_store->fetch({"id", "name"}, 100);
EXPECT_EQ(ipc_batch, nullptr);
}
TEST_F(MmapStoreTest, IPCFetchSingleRowWithInvalidColumn) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
ExecBatchPtr ipc_batch = ipc_store->fetch({"id", "invalid_column"}, 0);
EXPECT_EQ(ipc_batch, nullptr);
}
TEST_F(MmapStoreTest, IPCFetchSingleRowWithEmptyColumns) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
ExecBatchPtr ipc_batch = ipc_store->fetch({}, 0);
EXPECT_EQ(ipc_batch, nullptr);
}
TEST_F(MmapStoreTest, ParquetFetchSingleRowWithInvalidIndex) {
auto parquet_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(parquet_store->Open().ok());
ExecBatchPtr parquet_batch = parquet_store->fetch({"id", "name"}, -1);
EXPECT_EQ(parquet_batch, nullptr);
parquet_batch = parquet_store->fetch({"id", "name"}, 100);
EXPECT_EQ(parquet_batch, nullptr);
}
TEST_F(MmapStoreTest, AllDataType) {
auto mmap_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(mmap_store->Open().ok());
std::vector<std::string> columns = {"id", "list_int32"};
std::vector<int> indices = {0, 3, 6, 1, 0};
TablePtr mmap_table = mmap_store->fetch(columns, indices);
ASSERT_TRUE(mmap_table != nullptr);
EXPECT_EQ(mmap_table->num_rows(), 5);
EXPECT_EQ(mmap_table->num_columns(), 2);
for (size_t j = 0; j < columns.size(); ++j) {
auto column = mmap_table->column(j);
for (int k = 0; k < column->num_chunks(); ++k) {
auto array = column->chunk(k);
if (array->type()->id() == arrow::Type::INT32) {
auto int_array = std::static_pointer_cast<arrow::Int32Array>(array);
for (int i = 0; i < array->length(); ++i) {
int32_t value = int_array->Value(i);
EXPECT_EQ(value, indices[i] + 1);
}
} else if (array->type()->id() == arrow::Type::LIST) {
auto list_array = std::static_pointer_cast<arrow::ListArray>(array);
for (int i = 0; i < array->length(); ++i) {
auto list_value = list_array->value_slice(i);
auto list_value_array =
std::static_pointer_cast<arrow::Int32Array>(list_value);
EXPECT_EQ(list_value_array->length(), 128);
for (int m = 0; m < list_value_array->length(); ++m) {
int32_t value = list_value_array->Value(m);
EXPECT_EQ(value, indices[i] * 10 + m);
}
}
}
}
}
}
TEST_F(MmapStoreTest, FindRowGroupForRow) {
auto mmap_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(mmap_store->Open().ok());
EXPECT_EQ(mmap_store->FindRowGroupForRow(0), 0);
EXPECT_EQ(mmap_store->FindRowGroupForRow(1), 0);
EXPECT_EQ(mmap_store->FindRowGroupForRow(2), 0);
EXPECT_EQ(mmap_store->FindRowGroupForRow(3), 1);
EXPECT_EQ(mmap_store->FindRowGroupForRow(6), 2);
EXPECT_EQ(mmap_store->FindRowGroupForRow(9), 3);
EXPECT_EQ(mmap_store->FindRowGroupForRow(100), 3);
}
TEST_F(MmapStoreTest, GetRowGroupOffset) {
auto mmap_store = std::make_shared<MmapForwardStore>(parquet_path);
ASSERT_TRUE(mmap_store->Open().ok());
EXPECT_EQ(mmap_store->GetRowGroupOffset(0), 0);
EXPECT_EQ(mmap_store->GetRowGroupOffset(1), 3);
EXPECT_EQ(mmap_store->GetRowGroupOffset(2), 6);
EXPECT_EQ(mmap_store->GetRowGroupOffset(3), 9);
}
TEST_F(MmapStoreTest, InvalidPath) {
std::vector<std::string> err_path = {
"err_path",
"err_" + ipc_path,
"err_" + parquet_path,
ipc_path + ".unknown_file_type",
};
for (const auto &path : err_path) {
auto ipc_store = std::make_shared<MmapForwardStore>(path);
ASSERT_FALSE(ipc_store->Open().ok());
}
}
TEST_F(MmapStoreTest, InvalidFileFormat) {
std::string err_path = ipc_path + ".unknown_file_format";
EXPECT_EQ(InferFileFormat(err_path), FileFormat::UNKNOWN);
}
TEST_F(MmapStoreTest, ValidateEmptyColumns) {
auto ipc_store = std::make_shared<MmapForwardStore>(ipc_path);
ASSERT_TRUE(ipc_store->Open().ok());
EXPECT_FALSE(ipc_store->validate({}));
}
TEST_F(MmapStoreTest, ConstructorAndPhysicSchema) {
MmapForwardStore store(ipc_path);
EXPECT_EQ(store.physic_schema(), nullptr);
}
TEST_F(MmapStoreTest, DeleteDestructs) {
MmapForwardStore *store = new MmapForwardStore(ipc_path);
delete store;
}
@@ -0,0 +1,102 @@
// 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 "db/index/storage/parquet_writer.h"
#include <iostream>
#include <arrow/array/builder_primitive.h>
#include <arrow/record_batch.h>
#include <arrow/status.h>
#include <gtest/gtest.h>
using namespace zvec;
std::shared_ptr<arrow::RecordBatchReader> CreateTestReader(int start_id,
int count) {
auto schema = arrow::schema({arrow::field("id", arrow::int32()),
arrow::field("name", arrow::utf8())});
arrow::Int32Builder id_builder;
arrow::StringBuilder name_builder;
arrow::Status s;
for (int i = 0; i < count; ++i) {
s = id_builder.Append(start_id + i);
if (!s.ok()) {
return nullptr;
}
s = name_builder.Append("User" + std::to_string(start_id + i));
if (!s.ok()) {
return nullptr;
}
}
std::shared_ptr<arrow::Array> id_array, name_array;
s = id_builder.Finish(&id_array);
if (!s.ok()) {
return nullptr;
}
s = name_builder.Finish(&name_array);
if (!s.ok()) {
return nullptr;
}
auto batch = arrow::RecordBatch::Make(schema, count, {id_array, name_array});
auto maybe_reader = arrow::RecordBatchReader::Make({batch}, schema);
if (!maybe_reader.ok()) {
return nullptr;
}
return *maybe_reader;
}
TEST(ParquetWriter, General) {
ParquetWriter writer("output.parquet");
// writer.SetMaxRowsPerGroup(1000); // 可选:控制每组行数
// 第一次插入
{
auto reader1 = CreateTestReader(1, 3);
ASSERT_NE(reader1, nullptr);
auto status = writer.insert(reader1);
ASSERT_TRUE(status.ok());
std::cout << "Inserted batch 1" << std::endl;
}
// 第二次插入
{
auto reader2 = CreateTestReader(4, 2);
ASSERT_NE(reader2, nullptr);
auto status = writer.insert(reader2);
ASSERT_TRUE(status.ok());
std::cout << "Inserted batch 2" << std::endl;
}
// 第三次插入
{
auto reader3 = CreateTestReader(6, 4);
ASSERT_NE(reader3, nullptr);
auto status = writer.insert(reader3);
ASSERT_TRUE(status.ok());
std::cout << "Inserted batch 3" << std::endl;
}
// 最后关闭文件
auto status = writer.finalize();
if (!status.ok()) {
std::cerr << "Finalize failed: " << status.ToString() << std::endl;
}
std::cout << "Parquet file written successfully to output.parquet"
<< std::endl;
}
+737
View File
@@ -0,0 +1,737 @@
// 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.
#ifdef _MSC_VER
#define _ALLOW_KEYWORD_MACROS
#endif
#define private public
#define protected public
#include "db/index/storage/wal/wal_file.h"
#undef private
#undef protected
#include <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <stdio.h>
#include <gtest/gtest.h>
#include <zvec/ailego/parallel/thread_pool.h>
#include <zvec/ailego/utility/string_helper.h>
#include <zvec/ailego/utility/time_helper.h>
#include "db/common/file_helper.h"
#include "tests/test_util.h"
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
#endif
using namespace zvec;
using SegmentID = uint32_t;
class WalFileTest : public testing::Test {
protected:
void SetUp() override {
zvec::test_util::RemoveTestFiles("./data.wal.*");
}
void TearDown() override {}
};
TEST_F(WalFileTest, TestGeneral) {
std::string dir_path = "./";
SegmentID segment_id = 0;
std::string wal_file_path =
FileHelper::MakeFilePath(dir_path, FileID::WAL_FILE, segment_id);
WalFilePtr wal_file = WalFile::Create(wal_file_path);
ASSERT_TRUE(wal_file != nullptr);
WalOptions wal_option;
wal_option.create_new = true;
wal_option.max_docs_wal_flush = 0;
int ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
// add 100 same record
for (size_t i = 0; i < 100; i++) {
ret = wal_file->append(std::string("hello"));
ASSERT_EQ(ret, 0);
}
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// add 100-200 record
wal_option.create_new = false;
wal_option.max_docs_wal_flush = 1;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
for (size_t i = 100; i < 200; i++) {
std::string record = "hello";
ret = wal_file->append(record + std::to_string(i));
ASSERT_EQ(ret, 0);
}
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// reopen and add next 100 record
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
for (size_t i = 200; i < 300; i++) {
std::string record = "hello";
ret = wal_file->append(record + std::to_string(i));
ASSERT_EQ(ret, 0);
}
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// reopen and add batch model 100 record
wal_option.max_docs_wal_flush = 10;
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
for (size_t i = 300; i < 400; i++) {
std::string record = "hello";
ret = wal_file->append(record + std::to_string(i));
ASSERT_EQ(ret, 0);
}
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// reopen for read
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
uint32_t idx = 0;
ret = wal_file->prepare_for_read();
ASSERT_EQ(ret, 0);
std::string record = wal_file->next();
while (!record.empty()) {
if (idx < 100) {
ASSERT_EQ(record, "hello");
} else {
ASSERT_EQ(record, std::string("hello") + std::to_string(idx));
}
record = wal_file->next();
idx++;
}
ASSERT_EQ(idx, 400);
// close
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// remove
ret = wal_file->remove();
ASSERT_EQ(ret, 0);
}
void do_append(WalFile *wal_file, size_t number) {
std::string record = "hello" + std::to_string(number);
int ret = wal_file->append(std::move(record));
ASSERT_EQ(ret, 0);
}
TEST_F(WalFileTest, TestMultiThread) {
std::string dir_path = "./";
SegmentID segment_id = 0;
std::string wal_file_path =
FileHelper::MakeFilePath(dir_path, FileID::WAL_FILE, segment_id);
WalFilePtr wal_file = WalFile::Create(wal_file_path);
ASSERT_TRUE(wal_file != nullptr);
WalOptions wal_option;
wal_option.create_new = true;
wal_option.max_docs_wal_flush = 1;
int ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
ailego::ThreadPool pool(10, false);
for (size_t i = 0; i < 10000; i++) {
pool.execute(do_append, wal_file.get(), i);
}
pool.wait_finish();
wal_file->flush();
wal_file->close();
// reopen for batch model
wal_option.create_new = false;
wal_option.max_docs_wal_flush = 1000;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
for (size_t i = 10000; i < 20000; i++) {
pool.execute(do_append, wal_file.get(), i);
}
pool.wait_finish();
wal_file->flush();
wal_file->close();
// reopen for batch model
wal_option.create_new = false;
wal_option.max_docs_wal_flush = 0;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
for (size_t i = 20000; i < 30000; i++) {
pool.execute(do_append, wal_file.get(), i);
}
pool.wait_finish();
wal_file->flush();
wal_file->close();
// reopen for read
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
uint32_t idx = 0;
ret = wal_file->prepare_for_read();
ASSERT_EQ(ret, 0);
std::string record = wal_file->next();
while (!record.empty()) {
record = wal_file->next();
idx++;
}
ASSERT_EQ(idx, 30000);
// close
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// remove
ret = wal_file->remove();
ASSERT_EQ(ret, 0);
}
TEST_F(WalFileTest, TestBoundaryCondition) {
// read empty file
std::string dir_path = "./";
SegmentID segment_id = 0;
std::string wal_file_path =
FileHelper::MakeFilePath(dir_path, FileID::WAL_FILE, segment_id);
WalFilePtr wal_file = WalFile::Create(wal_file_path);
ASSERT_TRUE(wal_file != nullptr);
WalOptions wal_option;
wal_option.create_new = true;
wal_option.max_docs_wal_flush = 1;
int ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
uint32_t idx = 0;
std::string record = wal_file->next();
while (!record.empty()) {
record = wal_file->next();
idx++;
}
ASSERT_EQ(idx, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// write and read binary struct
std::vector<uint8_t> bin_v{0, 1, 2, 3};
std::string str(bin_v.begin(), bin_v.end());
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
ret = wal_file->append(std::move(str));
ASSERT_EQ(ret, 0);
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
idx = 0;
ret = wal_file->prepare_for_read();
ASSERT_EQ(ret, 0);
record = wal_file->next();
while (!record.empty()) {
ASSERT_EQ(record.size(), 4);
for (size_t i = 0; i < 4; i++) {
ASSERT_EQ(record[i], i);
}
record = wal_file->next();
idx++;
}
ASSERT_EQ(idx, 1);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// remove
ret = wal_file->remove();
ASSERT_EQ(ret, 0);
// write very large record 4Mb
size_t BIG_DATA_SIZE = 4 * 1024 * 1024;
std::vector<uint8_t> big_data(BIG_DATA_SIZE);
for (size_t i = 0; i < BIG_DATA_SIZE; i++) {
big_data[i] = i % 256;
}
str.clear();
str.assign((const char *)big_data.data(), BIG_DATA_SIZE);
wal_option.create_new = true;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
ret = wal_file->append(std::move(str));
ASSERT_EQ(ret, 0);
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
idx = 0;
ret = wal_file->prepare_for_read();
ASSERT_EQ(ret, 0);
record = wal_file->next();
while (!record.empty()) {
ASSERT_EQ(record.size(), BIG_DATA_SIZE);
for (size_t i = 0; i < BIG_DATA_SIZE; i++) {
ASSERT_EQ((uint8_t)record[i], i % 256);
}
record = wal_file->next();
idx++;
}
ASSERT_EQ(idx, 1);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// remove
ret = wal_file->remove();
ASSERT_EQ(ret, 0);
// batch model 100, just add 99 record and close
wal_option.max_docs_wal_flush = 100;
wal_option.create_new = true;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
for (size_t i = 0; i < 99; i++) {
std::string record = "hello";
ret = wal_file->append(record + std::to_string(i));
ASSERT_EQ(ret, 0);
}
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
idx = 0;
ret = wal_file->prepare_for_read();
ASSERT_EQ(ret, 0);
record = wal_file->next();
while (!record.empty()) {
ASSERT_EQ(record, std::string("hello") + std::to_string(idx));
record = wal_file->next();
idx++;
}
ASSERT_EQ(idx, 99);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// remove
ret = wal_file->remove();
ASSERT_EQ(ret, 0);
}
TEST_F(WalFileTest, TestNotExistErrorCase) {
std::string dir_path = "./";
SegmentID segment_id = 0;
std::string wal_file_path =
FileHelper::MakeFilePath(dir_path, FileID::WAL_FILE, segment_id);
WalFilePtr wal_file = WalFile::Create(wal_file_path);
// reopen for read
WalOptions wal_option;
wal_option.create_new = false;
int ret = wal_file->open(wal_option);
ASSERT_EQ(ret, -1);
}
TEST_F(WalFileTest, TestFirstErrorCase) {
std::string dir_path = "./";
SegmentID segment_id = 0;
std::string wal_file_path =
FileHelper::MakeFilePath(dir_path, FileID::WAL_FILE, segment_id);
WalFilePtr wal_file = WalFile::Create(wal_file_path);
ASSERT_TRUE(wal_file != nullptr);
WalOptions wal_option;
wal_option.create_new = true;
wal_option.max_docs_wal_flush = 1;
int ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
// add 10 same record
for (size_t i = 0; i < 10; i++) {
ret = wal_file->append(std::string("hello"));
ASSERT_EQ(ret, 0);
}
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
std::string wal_path = ailego::StringHelper::Concat(
dir_path, "data.wal.", std::to_string(segment_id));
int wal_fd = open(wal_path.c_str(), O_RDWR, 0644);
ASSERT_GT(wal_fd, 0);
// destory first record
lseek(wal_fd, 64 + 8, SEEK_SET);
// write err data
char buf[6] = "nihao";
write(wal_fd, buf, 5);
close(wal_fd);
// reopen for read
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
uint32_t idx = 0;
ret = wal_file->prepare_for_read();
ASSERT_EQ(ret, 0);
std::string record = wal_file->next();
while (!record.empty()) {
ASSERT_EQ(record, "hello");
record = wal_file->next();
idx++;
}
ASSERT_EQ(idx, 0);
// close
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// remove
ret = wal_file->remove();
ASSERT_EQ(ret, 0);
}
TEST_F(WalFileTest, TestMiddleErrorCase) {
std::string dir_path = "./";
SegmentID segment_id = 0;
std::string wal_file_path =
FileHelper::MakeFilePath(dir_path, FileID::WAL_FILE, segment_id);
WalFilePtr wal_file = WalFile::Create(wal_file_path);
ASSERT_TRUE(wal_file != nullptr);
WalOptions wal_option;
wal_option.create_new = true;
wal_option.max_docs_wal_flush = 1;
int ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
// add 10 same record
for (size_t i = 0; i < 10; i++) {
ret = wal_file->append(std::string("hello"));
ASSERT_EQ(ret, 0);
}
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
std::string wal_path = ailego::StringHelper::Concat(
dir_path, "data.wal.", std::to_string(segment_id));
int wal_fd = open(wal_path.c_str(), O_RDWR, 0644);
ASSERT_GT(wal_fd, 0);
// destory middle record
lseek(wal_fd, 64 + 13 * 5 + 8, SEEK_SET);
// write err data
char buf[6] = "nihao";
write(wal_fd, buf, 5);
close(wal_fd);
// reopen for read
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
uint32_t idx = 0;
ret = wal_file->prepare_for_read();
ASSERT_EQ(ret, 0);
std::string record = wal_file->next();
while (!record.empty()) {
ASSERT_EQ(record, "hello");
record = wal_file->next();
idx++;
}
ASSERT_EQ(idx, 5);
// close
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// remove
ret = wal_file->remove();
ASSERT_EQ(ret, 0);
}
TEST_F(WalFileTest, TestLastErrorCase) {
std::string dir_path = "./";
SegmentID segment_id = 0;
std::string wal_file_path =
FileHelper::MakeFilePath(dir_path, FileID::WAL_FILE, segment_id);
WalFilePtr wal_file = WalFile::Create(wal_file_path);
ASSERT_TRUE(wal_file != nullptr);
WalOptions wal_option;
wal_option.create_new = true;
wal_option.max_docs_wal_flush = 1;
int ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
// add 10 same record
for (size_t i = 0; i < 10; i++) {
ret = wal_file->append(std::string("hello"));
ASSERT_EQ(ret, 0);
}
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// destory last record
std::string wal_path = ailego::StringHelper::Concat(
dir_path, "data.wal.", std::to_string(segment_id));
int wal_fd = open(wal_path.c_str(), O_RDWR, 0644);
ASSERT_GT(wal_fd, 0);
off_t fsize = lseek(wal_fd, 0, SEEK_END);
close(wal_fd);
truncate(wal_path.c_str(), (fsize - 4));
// reopen for read
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
uint32_t idx = 0;
ret = wal_file->prepare_for_read();
ASSERT_EQ(ret, 0);
std::string record = wal_file->next();
while (!record.empty()) {
ASSERT_EQ(record, "hello");
record = wal_file->next();
idx++;
}
ASSERT_EQ(idx, 9);
// close
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// remove
ret = wal_file->remove();
ASSERT_EQ(ret, 0);
}
TEST_F(WalFileTest, TestLengthSmallErrorCase) {
std::string dir_path = "./";
SegmentID segment_id = 0;
std::string wal_file_path =
FileHelper::MakeFilePath(dir_path, FileID::WAL_FILE, segment_id);
WalFilePtr wal_file = WalFile::Create(wal_file_path);
ASSERT_TRUE(wal_file != nullptr);
WalOptions wal_option;
wal_option.create_new = true;
wal_option.max_docs_wal_flush = 1;
int ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
// add 10 same record
for (size_t i = 0; i < 10; i++) {
ret = wal_file->append(std::string("hello"));
ASSERT_EQ(ret, 0);
}
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// write error length
std::string wal_path = ailego::StringHelper::Concat(
dir_path, "data.wal.", std::to_string(segment_id));
int wal_fd = open(wal_path.c_str(), O_RDWR, 0644);
ASSERT_GT(wal_fd, 0);
uint32_t err_length = 2;
lseek(wal_fd, 64, SEEK_SET);
write(wal_fd, (const void *)&err_length, 4);
close(wal_fd);
// reopen for read
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
uint32_t idx = 0;
ret = wal_file->prepare_for_read();
ASSERT_EQ(ret, 0);
std::string record = wal_file->next();
while (!record.empty()) {
ASSERT_EQ(record, "hello");
record = wal_file->next();
idx++;
}
ASSERT_EQ(idx, 0);
// close
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// remove
ret = wal_file->remove();
ASSERT_EQ(ret, 0);
}
TEST_F(WalFileTest, TestLengthBigErrorCase) {
std::string dir_path = "./";
SegmentID segment_id = 0;
std::string wal_file_path =
FileHelper::MakeFilePath(dir_path, FileID::WAL_FILE, segment_id);
WalFilePtr wal_file = WalFile::Create(wal_file_path);
ASSERT_TRUE(wal_file != nullptr);
WalOptions wal_option;
wal_option.create_new = true;
wal_option.max_docs_wal_flush = 1;
int ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
// add 10 same record
for (size_t i = 0; i < 10; i++) {
ret = wal_file->append(std::string("hello"));
ASSERT_EQ(ret, 0);
}
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// write error length
std::string wal_path = ailego::StringHelper::Concat(
dir_path, "data.wal.", std::to_string(segment_id));
int wal_fd = open(wal_path.c_str(), O_RDWR, 0644);
ASSERT_GT(wal_fd, 0);
uint32_t err_length = 200; // exceed file size 130
lseek(wal_fd, 64, SEEK_SET);
write(wal_fd, (const void *)&err_length, 4);
close(wal_fd);
// reopen for read
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
uint32_t idx = 0;
ret = wal_file->prepare_for_read();
ASSERT_EQ(ret, 0);
std::string record = wal_file->next();
while (!record.empty()) {
ASSERT_EQ(record, "hello");
record = wal_file->next();
idx++;
}
ASSERT_EQ(idx, 0);
// close
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// remove
ret = wal_file->remove();
ASSERT_EQ(ret, 0);
}
TEST_F(WalFileTest, TestCRCErrorCase) {
std::string dir_path = "./";
SegmentID segment_id = 0;
std::string wal_file_path =
FileHelper::MakeFilePath(dir_path, FileID::WAL_FILE, segment_id);
WalFilePtr wal_file = WalFile::Create(wal_file_path);
ASSERT_TRUE(wal_file != nullptr);
WalOptions wal_option;
wal_option.create_new = true;
wal_option.max_docs_wal_flush = 1;
int ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
// add 10 same record
for (size_t i = 0; i < 10; i++) {
ret = wal_file->append(std::string("hello"));
ASSERT_EQ(ret, 0);
}
ret = wal_file->flush();
ASSERT_EQ(ret, 0);
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// write error crc
std::string wal_path = ailego::StringHelper::Concat(
dir_path, "data.wal.", std::to_string(segment_id));
int wal_fd = open(wal_path.c_str(), O_RDWR, 0644);
ASSERT_GT(wal_fd, 0);
// second record crc 64+(4+4+len(hello))+4
lseek(wal_fd, 64 + 17, SEEK_SET);
uint32_t err_crc = 123;
write(wal_fd, (const void *)&err_crc, 4);
close(wal_fd);
// reopen for read
wal_option.create_new = false;
ret = wal_file->open(wal_option);
ASSERT_EQ(ret, 0);
uint32_t idx = 0;
ret = wal_file->prepare_for_read();
ASSERT_EQ(ret, 0);
std::string record = wal_file->next();
while (!record.empty()) {
ASSERT_EQ(record, "hello");
record = wal_file->next();
idx++;
}
ASSERT_EQ(idx, 1);
// close
ret = wal_file->close();
ASSERT_EQ(ret, 0);
// remove
ret = wal_file->remove();
ASSERT_EQ(ret, 0);
}
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif
+804
View File
@@ -0,0 +1,804 @@
// 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 "utils.h"
#include <cstdint>
#include <memory>
#include <vector>
#include <zvec/ailego/logger/logger.h>
#include "zvec/db/collection.h"
#include "zvec/db/doc.h"
#include "zvec/db/index_params.h"
#include "zvec/db/schema.h"
#include "zvec/db/status.h"
#include "zvec/db/type.h"
using namespace zvec;
using namespace zvec::test;
CollectionSchema::Ptr TestHelper::CreateTempSchema() {
auto schema = std::make_shared<CollectionSchema>("demo");
schema->set_max_doc_count_per_segment(1000);
schema->add_field(std::make_shared<FieldSchema>(
"id", DataType::INT64, false, std::make_shared<InvertIndexParams>(true)));
schema->add_field(std::make_shared<FieldSchema>(
"name", DataType::STRING, false,
std::make_shared<InvertIndexParams>(false)));
schema->add_field(
std::make_shared<FieldSchema>("weight", DataType::FLOAT, true));
schema->add_field(std::make_shared<FieldSchema>(
"dense", DataType::VECTOR_FP32, 128, false,
std::make_shared<HnswIndexParams>(MetricType::IP)));
schema->add_field(std::make_shared<FieldSchema>(
"sparse", DataType::SPARSE_VECTOR_FP32, 0, false,
std::make_shared<HnswIndexParams>(MetricType::IP)));
return schema;
}
CollectionSchema::Ptr TestHelper::CreateScalarSchema() {
auto schema = std::make_shared<CollectionSchema>("demo");
// scalar
schema->add_field(std::make_shared<FieldSchema>("int32", DataType::INT32));
schema->add_field(std::make_shared<FieldSchema>("string", DataType::STRING));
return schema;
}
// Helper function
CollectionSchema::Ptr TestHelper::CreateNormalSchema(
bool nullable, std::string name, IndexParams::Ptr scalar_index_params,
IndexParams::Ptr vector_index_params, uint64_t max_doc_count) {
auto schema = std::make_shared<CollectionSchema>(name);
schema->set_max_doc_count_per_segment(max_doc_count);
// scalar
schema->add_field(std::make_shared<FieldSchema>(
"int32", DataType::INT32, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"string", DataType::STRING, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"uint32", DataType::UINT32, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"bool", DataType::BOOL, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"float", DataType::FLOAT, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"double", DataType::DOUBLE, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"int64", DataType::INT64, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"uint64", DataType::UINT64, nullable, scalar_index_params));
// array
schema->add_field(std::make_shared<FieldSchema>(
"array_int32", DataType::ARRAY_INT32, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"array_string", DataType::ARRAY_STRING, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"array_uint32", DataType::ARRAY_UINT32, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"array_bool", DataType::ARRAY_BOOL, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"array_float", DataType::ARRAY_FLOAT, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"array_double", DataType::ARRAY_DOUBLE, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"array_int64", DataType::ARRAY_INT64, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"array_uint64", DataType::ARRAY_UINT64, nullable, scalar_index_params));
schema->add_field(std::make_shared<FieldSchema>(
"dense_fp32", DataType::VECTOR_FP32, 128, false,
vector_index_params ? vector_index_params
: std::make_shared<FlatIndexParams>(MetricType::IP)));
schema->add_field(std::make_shared<FieldSchema>(
"dense_fp16", DataType::VECTOR_FP16, 128, false,
std::make_shared<FlatIndexParams>(MetricType::IP)));
schema->add_field(std::make_shared<FieldSchema>(
"dense_int8", DataType::VECTOR_INT8, 128, false,
std::make_shared<FlatIndexParams>(MetricType::IP)));
// IVF, HNSW_RABITQ and DISKANN do not support sparse vectors, always use
// Flat for sparse fields in those cases.
auto supports_sparse = [](const IndexParams::Ptr &params) {
auto type = params->type();
return type != IndexType::IVF && type != IndexType::HNSW_RABITQ &&
type != IndexType::DISKANN;
};
IndexParams::Ptr sparse_index_params;
if (vector_index_params && supports_sparse(vector_index_params)) {
sparse_index_params = vector_index_params->clone();
auto v = std::dynamic_pointer_cast<VectorIndexParams>(sparse_index_params);
// sparse always use IP
v->set_metric_type(MetricType::IP);
}
schema->add_field(std::make_shared<FieldSchema>(
"sparse_fp32", DataType::SPARSE_VECTOR_FP32, 128, false,
sparse_index_params ? sparse_index_params
: std::make_shared<FlatIndexParams>(MetricType::IP)));
schema->add_field(std::make_shared<FieldSchema>(
"sparse_fp16", DataType::SPARSE_VECTOR_FP16, 128, false,
std::make_shared<FlatIndexParams>(MetricType::IP)));
return schema;
}
CollectionSchema::Ptr TestHelper::CreateSchemaWithScalarIndex(
bool nullable, bool enable_optimize, std::string name) {
return CreateNormalSchema(
nullable, name, std::make_shared<InvertIndexParams>(enable_optimize));
}
CollectionSchema::Ptr TestHelper::CreateSchemaWithVectorIndex(
bool nullable, std::string name, IndexParams::Ptr vector_index_params) {
return CreateNormalSchema(
nullable, name, nullptr,
vector_index_params ? vector_index_params
: std::make_shared<HnswIndexParams>(MetricType::IP));
}
CollectionSchema::Ptr TestHelper::CreateSchemaWithMaxDocCount(
uint64_t doc_count) {
return CreateNormalSchema(false, "demo", nullptr, nullptr, doc_count);
}
std::string TestHelper::MakePK(const uint64_t doc_id) {
return "pk_" + std::to_string(doc_id);
}
uint64_t TestHelper::ExtractDocId(const std::string &pk) {
return std::stoull(pk.substr(3));
}
Doc TestHelper::CreateDoc(const uint64_t doc_id, const CollectionSchema &schema,
std::string pk) {
Doc new_doc;
if (pk.empty()) {
pk = MakePK(doc_id);
}
new_doc.set_pk(pk);
for (auto &field : schema.fields()) {
switch (field->data_type()) {
case DataType::BINARY: {
std::string binary_str("binary_" + std::to_string(doc_id));
new_doc.set<std::string>(field->name(), binary_str);
break;
}
case DataType::BOOL:
new_doc.set<bool>(field->name(), doc_id % 10 == 0);
break;
case DataType::INT32:
new_doc.set<int32_t>(field->name(), (int32_t)doc_id);
break;
case DataType::INT64:
new_doc.set<int64_t>(field->name(), (int64_t)doc_id);
break;
case DataType::UINT32:
new_doc.set<uint32_t>(field->name(), (uint32_t)doc_id);
break;
case DataType::UINT64:
new_doc.set<uint64_t>(field->name(), (uint64_t)doc_id);
break;
case DataType::FLOAT:
new_doc.set<float>(field->name(), (float)doc_id);
break;
case DataType::DOUBLE:
new_doc.set<double>(field->name(), (double)doc_id);
break;
case DataType::STRING:
new_doc.set<std::string>(field->name(),
"value_" + std::to_string(doc_id));
break;
case DataType::ARRAY_BINARY: {
std::vector<std::string> bin_vec;
for (size_t i = 0; i < (doc_id % 10); i++) {
bin_vec.push_back("bin_" + std::to_string(i));
}
new_doc.set<std::vector<std::string>>(field->name(), bin_vec);
break;
}
case DataType::ARRAY_BOOL:
new_doc.set<std::vector<bool>>(field->name(),
std::vector<bool>(10, doc_id % 10 == 0));
break;
case DataType::ARRAY_INT32:
new_doc.set<std::vector<int32_t>>(
field->name(), std::vector<int32_t>(10, (int32_t)doc_id));
break;
case DataType::ARRAY_INT64:
new_doc.set<std::vector<int64_t>>(
field->name(), std::vector<int64_t>(10, (int64_t)doc_id));
break;
case DataType::ARRAY_UINT32:
new_doc.set<std::vector<uint32_t>>(
field->name(), std::vector<uint32_t>(10, (uint32_t)doc_id));
break;
case DataType::ARRAY_UINT64:
new_doc.set<std::vector<uint64_t>>(
field->name(), std::vector<uint64_t>(10, (uint64_t)doc_id));
break;
case DataType::ARRAY_FLOAT:
new_doc.set<std::vector<float>>(field->name(),
std::vector<float>(10, (float)doc_id));
break;
case DataType::ARRAY_DOUBLE:
new_doc.set<std::vector<double>>(
field->name(), std::vector<double>(10, (double)doc_id));
break;
case DataType::ARRAY_STRING:
new_doc.set<std::vector<std::string>>(
field->name(),
std::vector<std::string>(10, "value_" + std::to_string(doc_id)));
break;
case DataType::VECTOR_BINARY32:
new_doc.set<std::vector<uint32_t>>(
field->name(),
std::vector<uint32_t>(field->dimension(), uint32_t(doc_id + 0.1)));
break;
case DataType::VECTOR_BINARY64:
new_doc.set<std::vector<uint64_t>>(
field->name(),
std::vector<uint64_t>(field->dimension(), uint64_t(doc_id + 0.1)));
break;
case DataType::VECTOR_FP32:
new_doc.set<std::vector<float>>(
field->name(),
std::vector<float>(field->dimension(), float(doc_id + 0.1)));
break;
case DataType::VECTOR_FP64:
new_doc.set<std::vector<double>>(
field->name(),
std::vector<double>(field->dimension(), double(doc_id + 0.1)));
break;
case DataType::VECTOR_FP16:
new_doc.set<std::vector<float16_t>>(
field->name(), std::vector<float16_t>(
field->dimension(),
static_cast<float16_t>(float(doc_id + 0.1))));
break;
case DataType::VECTOR_INT8:
new_doc.set<std::vector<int8_t>>(
field->name(),
std::vector<int8_t>(field->dimension(), (int8_t)doc_id));
break;
case DataType::VECTOR_INT16:
new_doc.set<std::vector<int16_t>>(
field->name(),
std::vector<int16_t>(field->dimension(), (int16_t)doc_id));
break;
case DataType::SPARSE_VECTOR_FP16: {
std::vector<uint32_t> indices;
std::vector<float16_t> values;
for (uint32_t i = 0; i < 100; i++) {
indices.push_back(i);
values.push_back(float16_t(float(doc_id + 0.1)));
}
std::pair<std::vector<uint32_t>, std::vector<float16_t>>
sparse_float_vec;
sparse_float_vec.first = indices;
sparse_float_vec.second = values;
new_doc.set<std::pair<std::vector<uint32_t>, std::vector<float16_t>>>(
field->name(), sparse_float_vec);
break;
}
case DataType::SPARSE_VECTOR_FP32: {
std::vector<uint32_t> indices;
std::vector<float> values;
for (uint32_t i = 0; i < 100; i++) {
indices.push_back(i);
values.push_back(float(doc_id + 0.1));
}
std::pair<std::vector<uint32_t>, std::vector<float>> sparse_float_vec;
sparse_float_vec.first = indices;
sparse_float_vec.second = values;
new_doc.set<std::pair<std::vector<uint32_t>, std::vector<float>>>(
field->name(), sparse_float_vec);
break;
}
default:
std::cout << "Unsupported data type: " << field->name() << std::endl;
throw std::runtime_error("Unsupported vector data type");
}
}
return new_doc;
}
Doc TestHelper::CreateDocNull(const uint64_t doc_id,
const CollectionSchema &schema, std::string pk) {
Doc new_doc;
if (pk.empty()) {
pk = "pk_" + std::to_string(doc_id);
}
new_doc.set_pk(pk);
for (auto &field : schema.fields()) {
switch (field->data_type()) {
case DataType::BINARY:
case DataType::BOOL:
case DataType::INT32:
case DataType::INT64:
case DataType::UINT32:
case DataType::UINT64:
case DataType::FLOAT:
case DataType::DOUBLE:
case DataType::STRING:
case DataType::ARRAY_BINARY:
case DataType::ARRAY_BOOL:
case DataType::ARRAY_INT32:
case DataType::ARRAY_INT64:
case DataType::ARRAY_UINT32:
case DataType::ARRAY_UINT64:
case DataType::ARRAY_FLOAT:
case DataType::ARRAY_DOUBLE:
case DataType::ARRAY_STRING:
break;
case DataType::VECTOR_FP32:
new_doc.set<std::vector<float>>(
field->name(),
std::vector<float>(field->dimension(), float(doc_id + 0.1)));
break;
case DataType::VECTOR_FP64:
new_doc.set<std::vector<double>>(
field->name(),
std::vector<double>(field->dimension(), double(doc_id + 0.1)));
break;
case DataType::VECTOR_FP16:
new_doc.set<std::vector<float16_t>>(
field->name(), std::vector<float16_t>(
field->dimension(),
static_cast<float16_t>(float(doc_id + 0.1))));
break;
case DataType::VECTOR_INT8:
new_doc.set<std::vector<int8_t>>(
field->name(),
std::vector<int8_t>(field->dimension(), (int8_t)doc_id));
break;
case DataType::VECTOR_INT16:
new_doc.set<std::vector<int16_t>>(
field->name(),
std::vector<int16_t>(field->dimension(), (int16_t)doc_id));
break;
case DataType::SPARSE_VECTOR_FP16: {
std::vector<uint32_t> indices;
std::vector<float16_t> values;
for (uint32_t i = 0; i < 100; i++) {
indices.push_back(i);
values.push_back(float16_t(float(doc_id + 0.1)));
}
std::pair<std::vector<uint32_t>, std::vector<float16_t>>
sparse_float_vec;
sparse_float_vec.first = indices;
sparse_float_vec.second = values;
new_doc.set<std::pair<std::vector<uint32_t>, std::vector<float16_t>>>(
field->name(), sparse_float_vec);
break;
}
case DataType::SPARSE_VECTOR_FP32: {
std::vector<uint32_t> indices;
std::vector<float> values;
for (uint32_t i = 0; i < 100; i++) {
indices.push_back(i);
values.push_back(float(doc_id + 0.1));
}
std::pair<std::vector<uint32_t>, std::vector<float>> sparse_float_vec;
sparse_float_vec.first = indices;
sparse_float_vec.second = values;
new_doc.set<std::pair<std::vector<uint32_t>, std::vector<float>>>(
field->name(), sparse_float_vec);
break;
}
default:
throw std::runtime_error("Unsupported vector data type");
}
}
return new_doc;
}
Status TestHelper::SegmentInsertDoc(const Segment::Ptr &segment,
const CollectionSchema &schema,
const uint64_t start_doc_id,
const uint64_t end_doc_id, bool nullable,
bool upsert, bool batch) {
for (auto doc_id = start_doc_id; doc_id < end_doc_id; doc_id++) {
if (segment) {
Doc new_doc;
if (nullable) {
new_doc = CreateDocNull(doc_id, schema);
} else {
new_doc = CreateDoc(doc_id, schema);
}
Status s;
if (upsert) {
s = segment->Upsert(new_doc);
CHECK_RETURN_STATUS(s);
} else {
s = segment->Insert(new_doc);
CHECK_RETURN_STATUS(s);
}
}
}
return Status::OK();
}
Status TestHelper::CollectionInsertDoc(const Collection::Ptr &collection,
const uint64_t start_doc_id,
const uint64_t end_doc_id, bool nullable,
bool upsert, bool batch) {
if (!collection) {
return Status::InvalidArgument("collection is nullptr");
}
auto schema = collection->Schema().value();
auto make_doc = [&](uint64_t doc_id) -> Doc {
return nullable ? CreateDocNull(doc_id, schema) : CreateDoc(doc_id, schema);
};
auto exec_write = [&](std::vector<Doc> &docs) -> Status {
Result<WriteResults> result =
upsert ? collection->Upsert(docs) : collection->Insert(docs);
if (!result.has_value()) {
LOG_ERROR("Failed to %s docs (count=%zu), error: %s.",
upsert ? "upsert" : "insert", docs.size(),
result.error().message().c_str());
return result.error();
}
const auto &write_results = result.value();
if (write_results.empty()) {
return Status::InternalError("WriteResults is unexpectedly empty");
}
for (const auto &wr : write_results) {
if (!wr.ok()) {
return wr;
}
}
return Status::OK();
};
if (batch) {
std::vector<Doc> docs;
docs.reserve(end_doc_id - start_doc_id);
for (uint64_t doc_id = start_doc_id; doc_id < end_doc_id; ++doc_id) {
docs.emplace_back(make_doc(doc_id));
}
return exec_write(docs);
} else {
std::vector<Doc> single_doc;
single_doc.reserve(1); // 可选优化
for (uint64_t doc_id = start_doc_id; doc_id < end_doc_id; ++doc_id) {
single_doc.clear();
single_doc.push_back(make_doc(doc_id));
Status s = exec_write(single_doc);
if (!s.ok()) {
LOG_ERROR("Failed at doc_id=%" PRIu64 ", doc: %s", doc_id,
single_doc[0].to_detail_string().c_str());
return s;
}
}
}
return Status::OK();
}
Status TestHelper::CollectionUpsertDoc(const Collection::Ptr &collection,
const uint64_t start_doc_id,
const uint64_t end_doc_id, bool nullable,
bool batch) {
return CollectionInsertDoc(collection, start_doc_id, end_doc_id, nullable,
true, batch);
}
Segment::Ptr TestHelper::CreateSegmentWithDoc(
const std::string &col_path, const CollectionSchema &schema,
SegmentID segment_id, uint64_t min_doc_id, const IDMap::Ptr &id_map,
const DeleteStore::Ptr &delete_store,
const VersionManager::Ptr &version_manager, const SegmentOptions &options,
uint64_t start_doc_id, uint32_t doc_count, bool nullable, bool upsert) {
auto result =
Segment::CreateAndOpen(col_path, schema, segment_id, min_doc_id, id_map,
delete_store, version_manager, options);
if (!result.has_value()) {
return nullptr;
}
auto segment = std::move(result).value();
auto s = SegmentInsertDoc(segment, schema, start_doc_id,
start_doc_id + doc_count, nullable, upsert);
if (!s.ok()) {
LOG_ERROR("Failed to insert doc, err: %s", s.message().c_str());
return nullptr;
}
return segment;
}
Collection::Ptr TestHelper::CreateCollectionWithDoc(
const std::string &path, const CollectionSchema &schema,
const CollectionOptions &options, uint64_t start_doc_id, uint32_t doc_count,
bool nullable, bool upsert) {
auto result = Collection::CreateAndOpen(path, schema, options);
if (!result.has_value()) {
LOG_ERROR("Failed to create collection, err: %s",
result.error().message().c_str());
return nullptr;
}
auto collection = std::move(result).value();
auto s = CollectionInsertDoc(collection, start_doc_id,
start_doc_id + doc_count, nullable, upsert);
if (!s.ok()) {
LOG_ERROR("Failed to insert doc, err: %s", s.message().c_str());
return nullptr;
}
return collection;
}
arrow::Status TestHelper::WriteTestFile(const std::string &filepath,
FileFormat format,
uint32_t start_doc_id,
uint32_t end_doc_id,
uint32_t batch_size) {
// Define schema with additional list types
auto schema = arrow::schema(
{arrow::field(GLOBAL_DOC_ID, arrow::uint64()),
arrow::field(USER_ID, arrow::utf8()), arrow::field("id", arrow::int32()),
arrow::field("name", arrow::utf8()),
arrow::field("score", arrow::float64()),
arrow::field("list_binary", arrow::list(arrow::binary())),
arrow::field("list_utf8", arrow::list(arrow::utf8())),
arrow::field("list_boolean", arrow::list(arrow::boolean())),
arrow::field("list_int32", arrow::list(arrow::int32())),
arrow::field("list_int64", arrow::list(arrow::int64())),
arrow::field("list_uint32", arrow::list(arrow::uint32())),
arrow::field("list_uint64", arrow::list(arrow::uint64())),
arrow::field("list_float32", arrow::list(arrow::float32())),
arrow::field("list_float64", arrow::list(arrow::float64()))});
// Create builders
auto g_doc_id_builder = std::make_shared<arrow::UInt64Builder>();
auto uid_builder = std::make_shared<arrow::StringBuilder>();
auto id_builder = std::make_shared<arrow::Int32Builder>();
auto name_builder = std::make_shared<arrow::StringBuilder>();
auto score_builder = std::make_shared<arrow::DoubleBuilder>();
// Array field builders
auto list_binary_builder = std::make_shared<arrow::ListBuilder>(
arrow::default_memory_pool(), std::make_shared<arrow::BinaryBuilder>());
auto list_utf8_builder = std::make_shared<arrow::ListBuilder>(
arrow::default_memory_pool(), std::make_shared<arrow::StringBuilder>());
auto list_boolean_builder = std::make_shared<arrow::ListBuilder>(
arrow::default_memory_pool(), std::make_shared<arrow::BooleanBuilder>());
auto list_int32_builder = std::make_shared<arrow::ListBuilder>(
arrow::default_memory_pool(), std::make_shared<arrow::Int32Builder>());
auto list_int64_builder = std::make_shared<arrow::ListBuilder>(
arrow::default_memory_pool(), std::make_shared<arrow::Int64Builder>());
auto list_uint32_builder = std::make_shared<arrow::ListBuilder>(
arrow::default_memory_pool(), std::make_shared<arrow::UInt32Builder>());
auto list_uint64_builder = std::make_shared<arrow::ListBuilder>(
arrow::default_memory_pool(), std::make_shared<arrow::UInt64Builder>());
auto list_float32_builder = std::make_shared<arrow::ListBuilder>(
arrow::default_memory_pool(), std::make_shared<arrow::FloatBuilder>());
auto list_float64_builder = std::make_shared<arrow::ListBuilder>(
arrow::default_memory_pool(), std::make_shared<arrow::DoubleBuilder>());
// Cast child builders for easier access
auto binary_builder =
static_cast<arrow::BinaryBuilder *>(list_binary_builder->value_builder());
auto utf8_child_builder =
static_cast<arrow::StringBuilder *>(list_utf8_builder->value_builder());
auto boolean_child_builder = static_cast<arrow::BooleanBuilder *>(
list_boolean_builder->value_builder());
auto int32_child_builder =
static_cast<arrow::Int32Builder *>(list_int32_builder->value_builder());
auto int64_child_builder =
static_cast<arrow::Int64Builder *>(list_int64_builder->value_builder());
auto uint32_child_builder =
static_cast<arrow::UInt32Builder *>(list_uint32_builder->value_builder());
auto uint64_child_builder =
static_cast<arrow::UInt64Builder *>(list_uint64_builder->value_builder());
auto float32_child_builder =
static_cast<arrow::FloatBuilder *>(list_float32_builder->value_builder());
auto float64_child_builder = static_cast<arrow::DoubleBuilder *>(
list_float64_builder->value_builder());
// Fill data
for (uint32_t i = start_doc_id; i < end_doc_id; ++i) {
ARROW_RETURN_NOT_OK(g_doc_id_builder->Append(i + 1));
ARROW_RETURN_NOT_OK(uid_builder->Append("user_" + std::to_string(i + 1)));
ARROW_RETURN_NOT_OK(id_builder->Append(i + 1));
ARROW_RETURN_NOT_OK(name_builder->Append("Name" + std::to_string(i)));
ARROW_RETURN_NOT_OK(score_builder->Append(80.0 + i));
const int dim = 128;
// Append list_binary data
ARROW_RETURN_NOT_OK(list_binary_builder->Append());
for (int j = 0; j < dim; ++j) {
std::string binary_data =
"binary_" + std::to_string(i) + "_" + std::to_string(j);
ARROW_RETURN_NOT_OK(binary_builder->Append(binary_data));
}
// Append list_utf8 data
ARROW_RETURN_NOT_OK(list_utf8_builder->Append());
for (int j = 0; j < dim; ++j) {
ARROW_RETURN_NOT_OK(utf8_child_builder->Append(
"string_" + std::to_string(i) + "_" + std::to_string(j)));
}
// Append list_boolean data
ARROW_RETURN_NOT_OK(list_boolean_builder->Append());
for (int j = 0; j < dim; ++j) {
ARROW_RETURN_NOT_OK(boolean_child_builder->Append((i + j) % 2 == 0));
}
// Append list_int32 data
ARROW_RETURN_NOT_OK(list_int32_builder->Append());
for (int j = 0; j < dim; ++j) {
ARROW_RETURN_NOT_OK(int32_child_builder->Append(i * 10 + j));
}
// Append list_int64 data
ARROW_RETURN_NOT_OK(list_int64_builder->Append());
for (int j = 0; j < dim; ++j) {
ARROW_RETURN_NOT_OK(
int64_child_builder->Append(static_cast<int64_t>(i) * 100 + j));
}
// Append list_uint32 data
ARROW_RETURN_NOT_OK(list_uint32_builder->Append());
for (int j = 0; j < dim; ++j) {
ARROW_RETURN_NOT_OK(
uint32_child_builder->Append(static_cast<uint32_t>(i) * 10 + j));
}
// Append list_uint64 data
ARROW_RETURN_NOT_OK(list_uint64_builder->Append());
for (int j = 0; j < dim; ++j) {
ARROW_RETURN_NOT_OK(
uint64_child_builder->Append(static_cast<uint64_t>(i) * 100 + j));
}
// Append list_float32 data
ARROW_RETURN_NOT_OK(list_float32_builder->Append());
for (int j = 0; j < dim; ++j) {
ARROW_RETURN_NOT_OK(
float32_child_builder->Append(static_cast<float>(i) + j * 0.1f));
}
// Append list_float64 data
ARROW_RETURN_NOT_OK(list_float64_builder->Append());
for (int j = 0; j < dim; ++j) {
ARROW_RETURN_NOT_OK(
float64_child_builder->Append(static_cast<double>(i) + j * 0.01));
}
}
// Construct arrays
std::shared_ptr<arrow::Array> g_doc_id_array, uid_array, id_array, name_array,
score_array, list_binary_array, list_utf8_array, list_boolean_array,
list_int32_array, list_int64_array, list_uint32_array, list_uint64_array,
list_float32_array, list_float64_array;
ARROW_RETURN_NOT_OK(g_doc_id_builder->Finish(&g_doc_id_array));
ARROW_RETURN_NOT_OK(uid_builder->Finish(&uid_array));
ARROW_RETURN_NOT_OK(id_builder->Finish(&id_array));
ARROW_RETURN_NOT_OK(name_builder->Finish(&name_array));
ARROW_RETURN_NOT_OK(score_builder->Finish(&score_array));
ARROW_RETURN_NOT_OK(list_binary_builder->Finish(&list_binary_array));
ARROW_RETURN_NOT_OK(list_utf8_builder->Finish(&list_utf8_array));
ARROW_RETURN_NOT_OK(list_boolean_builder->Finish(&list_boolean_array));
ARROW_RETURN_NOT_OK(list_int32_builder->Finish(&list_int32_array));
ARROW_RETURN_NOT_OK(list_int64_builder->Finish(&list_int64_array));
ARROW_RETURN_NOT_OK(list_uint32_builder->Finish(&list_uint32_array));
ARROW_RETURN_NOT_OK(list_uint64_builder->Finish(&list_uint64_array));
ARROW_RETURN_NOT_OK(list_float32_builder->Finish(&list_float32_array));
ARROW_RETURN_NOT_OK(list_float64_builder->Finish(&list_float64_array));
// Set rows per batch
std::vector<std::shared_ptr<arrow::RecordBatch>> batches;
// Split data into multiple batches
auto doc_count = (int)(end_doc_id - start_doc_id);
for (int start = 0; start < doc_count; start += batch_size) {
int current_batch_size = std::min((int)batch_size, doc_count - start);
auto g_doc_id_slice = g_doc_id_array->Slice(start, current_batch_size);
auto uid_slice = uid_array->Slice(start, current_batch_size);
auto id_slice = id_array->Slice(start, current_batch_size);
auto name_slice = name_array->Slice(start, current_batch_size);
auto score_slice = score_array->Slice(start, current_batch_size);
auto list_binary_slice =
list_binary_array->Slice(start, current_batch_size);
auto list_utf8_slice = list_utf8_array->Slice(start, current_batch_size);
auto list_boolean_slice =
list_boolean_array->Slice(start, current_batch_size);
auto list_int32_slice = list_int32_array->Slice(start, current_batch_size);
auto list_int64_slice = list_int64_array->Slice(start, current_batch_size);
auto list_uint32_slice =
list_uint32_array->Slice(start, current_batch_size);
auto list_uint64_slice =
list_uint64_array->Slice(start, current_batch_size);
auto list_float32_slice =
list_float32_array->Slice(start, current_batch_size);
auto list_float64_slice =
list_float64_array->Slice(start, current_batch_size);
auto batch = arrow::RecordBatch::Make(
schema, current_batch_size,
{g_doc_id_slice, uid_slice, id_slice, name_slice, score_slice,
list_binary_slice, list_utf8_slice, list_boolean_slice,
list_int32_slice, list_int64_slice, list_uint32_slice,
list_uint64_slice, list_float32_slice, list_float64_slice});
batches.push_back(batch);
}
// Open output stream
ARROW_ASSIGN_OR_RAISE(auto out, arrow::io::FileOutputStream::Open(filepath));
if (format == FileFormat::PARQUET) {
// Parquet write logic - create table with multiple record batches
auto table = arrow::Table::Make(
schema, {g_doc_id_array, uid_array, id_array, name_array, score_array,
list_binary_array, list_utf8_array, list_boolean_array,
list_int32_array, list_int64_array, list_uint32_array,
list_uint64_array, list_float32_array, list_float64_array});
parquet::WriterProperties::Builder builder;
builder.data_pagesize(1024);
// 3 rows per row group
builder.max_row_group_length(batch_size);
auto props = builder.build();
auto status = parquet::arrow::WriteTable(
*table, arrow::default_memory_pool(), out, batch_size, props);
if (!status.ok()) {
std::cerr << "Write failed: " << status.ToString() << std::endl;
return status;
}
std::cout << "Wrote test Parquet file with multiple row groups: "
<< filepath << std::endl;
} else if (format == FileFormat::IPC) {
// IPC write logic - write multiple record batches
auto writer_result = arrow::ipc::MakeFileWriter(out, schema);
ARROW_RETURN_NOT_OK(writer_result.status());
auto writer = std::move(writer_result).ValueOrDie();
// Write multiple batches
for (const auto &batch : batches) {
ARROW_RETURN_NOT_OK(writer->WriteRecordBatch(*batch));
}
ARROW_RETURN_NOT_OK(writer->Close());
std::cout << "Wrote test IPC file with " << batches.size()
<< " batches: " << filepath << std::endl;
}
ARROW_RETURN_NOT_OK(out->Close());
return arrow::Status::OK();
}
+219
View File
@@ -0,0 +1,219 @@
// 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.
#pragma once
#include <cstdint>
#include <iostream>
#include <memory>
#include <string>
#include <arrow/array/array_binary.h>
#include <arrow/io/file.h>
#include <arrow/ipc/reader.h>
#include <arrow/ipc/writer.h>
#include <arrow/pretty_print.h>
#include <arrow/result.h>
#include <arrow/table.h>
#include <gtest/gtest.h>
#include <parquet/arrow/writer.h>
#include "db/common/constants.h"
#include "db/common/typedef.h"
#include "db/index/common/meta.h"
#include "db/index/segment/segment.h"
#include "db/index/storage/store_helper.h"
#include "zvec/db/collection.h"
#include "zvec/db/doc.h"
#include "zvec/db/schema.h"
#include "zvec/db/type.h"
namespace zvec::test {
template <typename T>
bool vectors_equal_when_sorted(std::vector<T> a, std::vector<T> b) {
if (a.size() != b.size()) {
return false;
}
std::sort(a.begin(), a.end());
std::sort(b.begin(), b.end());
return a == b;
}
template <typename T>
double inner_produce_double(const std::vector<T> &vec1,
const std::vector<T> &vec2) {
double result = 0.0;
for (size_t i = 0; i < vec1.size(); ++i) {
result += vec1[i] * vec2[i];
}
return result;
}
template <typename T>
inline float cosine_distance_dense(const std::vector<T> &vec1,
const std::vector<T> &vec2) {
const auto dot = inner_produce_double(vec1, vec2);
const auto norm1 = std::sqrt((inner_produce_double(vec1, vec1)));
const auto norm2 = std::sqrt((inner_produce_double(vec2, vec2)));
if (norm1 == 0.0f || norm2 == 0.0f) return 0.0f;
return 1.0f - dot / (norm1 * norm2);
}
template <typename T>
inline float dp_distance_dense(const std::vector<T> &vec1,
const std::vector<T> &vec2) {
double result = 0.0;
for (size_t i = 0; i < vec1.size(); ++i) {
result += vec1[i] * vec2[i];
}
return result;
}
template <typename T>
inline float euclidean_distance_dense(const std::vector<T> &vec1,
const std::vector<T> &vec2) {
double sum = 0.0f;
for (size_t i = 0; i < vec1.size(); ++i) {
const float diff =
static_cast<float>(vec1[i]) - static_cast<float>(vec2[i]);
sum += diff * diff;
}
return sum;
}
template <typename T>
inline float distance_dense(const std::vector<T> &vec1,
const std::vector<T> &vec2, MetricType metric) {
switch (metric) {
case MetricType::COSINE:
return cosine_distance_dense(vec1, vec2);
case MetricType::L2:
return euclidean_distance_dense(vec1, vec2);
case MetricType::IP:
return dp_distance_dense(vec1, vec2);
default:
throw std::invalid_argument("Unsupported metric for FP32");
}
}
using SparseVecFP32 = std::pair<std::vector<uint32_t>, std::vector<float>>;
using SparseVecFP16 = std::pair<std::vector<uint32_t>, std::vector<float16_t>>;
using SparseVec = SparseVecFP32;
template <typename T>
inline float sparse_dot_product(const std::vector<uint32_t> &idx1,
const std::vector<T> &val1,
const std::vector<uint32_t> &idx2,
const std::vector<T> &val2) {
double dot = 0.0f;
size_t i = 0, j = 0;
while (i < idx1.size() && j < idx2.size()) {
if (idx1[i] == idx2[j]) {
dot += static_cast<float>(val1[i]) * static_cast<float>(val2[j]);
++i;
++j;
} else if (idx1[i] < idx2[j]) {
++i;
} else {
++j;
}
}
return dot;
}
inline float distance_sparse(const SparseVecFP32 &vec1,
const SparseVecFP32 &vec2) {
return sparse_dot_product(vec1.first, vec1.second, vec2.first, vec2.second);
}
inline float distance_sparse(const SparseVecFP16 &vec1,
const SparseVecFP16 &vec2) {
return sparse_dot_product(vec1.first, vec1.second, vec2.first, vec2.second);
}
class TestHelper {
public:
static CollectionSchema::Ptr CreateTempSchema();
static CollectionSchema::Ptr CreateScalarSchema();
static CollectionSchema::Ptr CreateNormalSchema(
bool nullable = false, std::string name = "demo",
IndexParams::Ptr scalar_index_params = nullptr,
IndexParams::Ptr vector_index_params = nullptr,
uint64_t max_doc_count = MAX_DOC_COUNT_PER_SEGMENT);
static CollectionSchema::Ptr CreateSchemaWithScalarIndex(
bool nullable = false, bool enable_optimize = false,
std::string name = "demo");
static CollectionSchema::Ptr CreateSchemaWithVectorIndex(
bool nullable = false, std::string name = "demo",
IndexParams::Ptr vector_index_params = nullptr);
static CollectionSchema::Ptr CreateSchemaWithMaxDocCount(uint64_t doc_count);
static std::string MakePK(const uint64_t doc_id);
static uint64_t ExtractDocId(const std::string &pk);
static Doc CreateDoc(const uint64_t doc_id, const CollectionSchema &schema,
std::string pk = "");
static Doc CreateDocNull(const uint64_t doc_id,
const CollectionSchema &schema, std::string pk = "");
static Status SegmentInsertDoc(const Segment::Ptr &segment,
const CollectionSchema &schema,
const uint64_t start_doc_id,
const uint64_t end_doc_id,
bool nullable = false, bool upsert = false,
bool batch = false);
static Status CollectionInsertDoc(const Collection::Ptr &collection,
const uint64_t start_doc_id,
const uint64_t end_doc_id,
bool nullable = false, bool upsert = false,
bool batch = false);
static Status CollectionUpsertDoc(const Collection::Ptr &collection,
const uint64_t start_doc_id,
const uint64_t end_doc_id,
bool nullable = false, bool batch = false);
static Segment::Ptr CreateSegmentWithDoc(
const std::string &col_path, const CollectionSchema &schema,
SegmentID segment_id, uint64_t min_doc_id, const IDMap::Ptr &id_map,
const DeleteStore::Ptr &delete_store,
const VersionManager::Ptr &version_manager, const SegmentOptions &options,
uint64_t start_doc_id, uint32_t doc_count, bool nullable = false,
bool upsert = false);
static Collection::Ptr CreateCollectionWithDoc(
const std::string &path, const CollectionSchema &schema,
const CollectionOptions &options, uint64_t start_doc_id,
uint32_t doc_count, bool nullable = false, bool upsert = false);
static arrow::Status WriteTestFile(const std::string &filepath,
FileFormat format,
uint32_t start_doc_id = 0,
uint32_t end_doc_id = 10,
uint32_t batch_size = 3);
};
} // namespace zvec::test
+315
View File
@@ -0,0 +1,315 @@
// 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.
#define _USE_MATH_DEFINES
#include <algorithm>
#include <cmath>
#include <memory>
#include <set>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include <zvec/db/doc.h>
#include <zvec/db/index_params.h>
#include <zvec/db/reranker.h>
#include <zvec/db/type.h>
using namespace zvec;
namespace {
Doc::Ptr MakeDoc(const std::string &id, float score) {
auto doc = std::make_shared<Doc>();
doc->set_pk(id);
doc->set_score(score);
return doc;
}
FieldSchema::Ptr MakeField(const std::string &name, MetricType metric) {
return std::make_shared<FieldSchema>(
name, DataType::VECTOR_FP16, /*dimension=*/4, /*nullable=*/false,
std::make_shared<HnswIndexParams>(metric));
}
} // namespace
// ==================== RRF Tests ====================
TEST(RerankRrfTest, BasicRRF) {
// Two sub-queries, each returning 3 documents with some overlap.
std::vector<DocPtrList> results;
results.push_back(
{MakeDoc("a", 0.9f), MakeDoc("b", 0.8f), MakeDoc("c", 0.7f)});
results.push_back(
{MakeDoc("b", 0.95f), MakeDoc("a", 0.85f), MakeDoc("d", 0.75f)});
auto result =
reranker::rerank(reranker::RrfParams{/*rank_constant=*/60}, results,
/*fields=*/{}, /*topn=*/10);
ASSERT_TRUE(result.has_value());
auto &out = result.value();
// "a" appears at rank 0 in sub-query 0 and rank 1 in sub-query 1:
// rrf_score = 1/(60+0+1) + 1/(60+1+1) = 1/61 + 1/62
// "b" appears at rank 1 in sub-query 0 and rank 0 in sub-query 1:
// rrf_score = 1/(60+1+1) + 1/(60+0+1) = 1/62 + 1/61
// So a and b should have equal scores and occupy the top two slots.
ASSERT_GE(out.size(), 3u);
std::set<std::string> top2 = {out[0]->pk(), out[1]->pk()};
EXPECT_EQ(top2, (std::set<std::string>{"a", "b"}));
EXPECT_NEAR(out[0]->score(), out[1]->score(), 1e-10);
}
TEST(RerankRrfTest, Topn) {
std::vector<DocPtrList> results;
results.push_back(
{MakeDoc("a", 0.9f), MakeDoc("b", 0.8f), MakeDoc("c", 0.7f)});
auto result =
reranker::rerank(reranker::RrfParams{/*rank_constant=*/60}, results,
/*fields=*/{}, /*topn=*/2);
ASSERT_TRUE(result.has_value());
ASSERT_EQ(result.value().size(), 2u);
}
TEST(RerankRrfTest, SingleField) {
std::vector<DocPtrList> results;
results.push_back({MakeDoc("a", 0.9f), MakeDoc("b", 0.8f)});
auto result =
reranker::rerank(reranker::RrfParams{/*rank_constant=*/60}, results,
/*fields=*/{}, /*topn=*/10);
ASSERT_TRUE(result.has_value());
auto &out = result.value();
ASSERT_EQ(out.size(), 2u);
// With single sub-query, RRF score for rank 0 > rank 1.
EXPECT_GT(out[0]->score(), out[1]->score());
}
TEST(RerankRrfTest, EmptyResults) {
std::vector<DocPtrList> results;
auto result =
reranker::rerank(reranker::RrfParams{/*rank_constant=*/60}, results,
/*fields=*/{}, /*topn=*/10);
ASSERT_TRUE(result.has_value());
EXPECT_TRUE(result.value().empty());
}
TEST(RerankRrfTest, DefaultParams) {
// RrfParams (and therefore RerankParams) defaults to rank_constant = 60.
std::vector<DocPtrList> results;
results.push_back({MakeDoc("a", 0.9f), MakeDoc("b", 0.8f)});
auto result =
reranker::rerank(reranker::RerankParams{}, results, /*fields=*/{},
/*topn=*/10);
ASSERT_TRUE(result.has_value());
ASSERT_EQ(result.value().size(), 2u);
}
// ==================== Weighted Tests ====================
TEST(RerankWeightedTest, BasicWeighted) {
std::vector<DocPtrList> results;
results.push_back({MakeDoc("a", 0.5f), MakeDoc("b", 0.3f)});
results.push_back({MakeDoc("a", 0.8f), MakeDoc("c", 0.6f)});
std::vector<FieldSchema::Ptr> fields = {MakeField("vec1", MetricType::L2),
MakeField("vec2", MetricType::L2)};
auto result =
reranker::rerank(reranker::WeightedParams{{0.7, 0.3}}, results, fields,
/*topn=*/10);
ASSERT_TRUE(result.has_value());
auto &out = result.value();
ASSERT_GE(out.size(), 2u);
// "a" appears in both sub-queries, should have highest combined score.
EXPECT_EQ(out[0]->pk(), "a");
}
TEST(RerankWeightedTest, MixedMetrics) {
std::vector<DocPtrList> results;
results.push_back({MakeDoc("a", 0.5f)});
results.push_back({MakeDoc("a", 0.4f)});
std::vector<FieldSchema::Ptr> fields = {
MakeField("vec1", MetricType::L2), MakeField("vec2", MetricType::COSINE)};
auto result =
reranker::rerank(reranker::WeightedParams{{0.5, 0.5}}, results, fields,
/*topn=*/10);
ASSERT_TRUE(result.has_value());
auto &out = result.value();
ASSERT_EQ(out.size(), 1u);
EXPECT_EQ(out[0]->pk(), "a");
// L2 normalize(0.5) = 1 - 2*atan(0.5)/pi
// COSINE normalize(0.4) = 1 - 0.4/2 = 0.8
// weighted = l2_norm * 0.5 + cos_norm * 0.5
double l2_norm = 1.0 - 2.0 * std::atan(0.5) / M_PI;
double cos_norm = 1.0 - 0.4 / 2.0;
double expected = l2_norm * 0.5 + cos_norm * 0.5;
EXPECT_NEAR(out[0]->score(), expected, 1e-5);
}
TEST(RerankWeightedTest, WeightsCountMismatch) {
std::vector<DocPtrList> results;
results.push_back({MakeDoc("a", 0.5f)});
results.push_back({MakeDoc("b", 0.3f)});
std::vector<FieldSchema::Ptr> fields = {MakeField("vec1", MetricType::L2),
MakeField("vec2", MetricType::L2)};
// Only one weight provided for two sub-queries.
auto result =
reranker::rerank(reranker::WeightedParams{{1.0}}, results, fields,
/*topn=*/10);
ASSERT_FALSE(result.has_value());
}
TEST(RerankWeightedTest, FieldsCountMismatch) {
std::vector<DocPtrList> results;
results.push_back({MakeDoc("a", 0.5f)});
results.push_back({MakeDoc("b", 0.3f)});
std::vector<FieldSchema::Ptr> fields = {MakeField("vec1", MetricType::L2)};
auto result =
reranker::rerank(reranker::WeightedParams{{0.5, 0.5}}, results, fields,
/*topn=*/10);
ASSERT_FALSE(result.has_value());
}
TEST(RerankWeightedTest, NullFieldError) {
std::vector<DocPtrList> results;
results.push_back({MakeDoc("a", 0.5f)});
std::vector<FieldSchema::Ptr> fields = {nullptr};
auto result =
reranker::rerank(reranker::WeightedParams{{1.0}}, results, fields,
/*topn=*/10);
ASSERT_FALSE(result.has_value());
}
TEST(RerankWeightedTest, NormalizeL2) {
std::vector<DocPtrList> results;
results.push_back({MakeDoc("a", 0.0f), MakeDoc("b", 1.0f)});
std::vector<FieldSchema::Ptr> fields = {MakeField("vec1", MetricType::L2)};
auto result =
reranker::rerank(reranker::WeightedParams{{1.0}}, results, fields,
/*topn=*/10);
ASSERT_TRUE(result.has_value());
auto &out = result.value();
ASSERT_EQ(out.size(), 2u);
// L2 normalize(0.0) = 1.0, normalize(1.0) in (0, 1)
EXPECT_NEAR(out[0]->score(), 1.0, 1e-10);
EXPECT_EQ(out[0]->pk(), "a");
EXPECT_GT(out[1]->score(), 0.0);
EXPECT_LT(out[1]->score(), 1.0);
}
TEST(RerankWeightedTest, NormalizeIP) {
std::vector<DocPtrList> results;
results.push_back({MakeDoc("a", 0.0f), MakeDoc("b", 1.0f)});
std::vector<FieldSchema::Ptr> fields = {MakeField("vec1", MetricType::IP)};
auto result =
reranker::rerank(reranker::WeightedParams{{1.0}}, results, fields,
/*topn=*/10);
ASSERT_TRUE(result.has_value());
auto &out = result.value();
ASSERT_EQ(out.size(), 2u);
// IP normalize(1.0) > 0.5 > normalize(0.0) = 0.5
EXPECT_EQ(out[0]->pk(), "b");
EXPECT_GT(out[0]->score(), 0.5);
EXPECT_NEAR(out[1]->score(), 0.5, 1e-10);
}
TEST(RerankWeightedTest, NormalizeCosine) {
std::vector<DocPtrList> results;
results.push_back(
{MakeDoc("a", 0.0f), MakeDoc("b", 1.0f), MakeDoc("c", 2.0f)});
std::vector<FieldSchema::Ptr> fields = {
MakeField("vec1", MetricType::COSINE)};
auto result =
reranker::rerank(reranker::WeightedParams{{1.0}}, results, fields,
/*topn=*/10);
ASSERT_TRUE(result.has_value());
auto &out = result.value();
ASSERT_EQ(out.size(), 3u);
// COSINE normalize(0.0) = 1.0, normalize(1.0) = 0.5, normalize(2.0) = 0.0
EXPECT_NEAR(out[0]->score(), 1.0, 1e-10);
EXPECT_NEAR(out[1]->score(), 0.5, 1e-10);
EXPECT_NEAR(out[2]->score(), 0.0, 1e-10);
}
TEST(RerankWeightedTest, Topn) {
std::vector<DocPtrList> results;
results.push_back(
{MakeDoc("a", 0.1f), MakeDoc("b", 0.2f), MakeDoc("c", 0.3f)});
std::vector<FieldSchema::Ptr> fields = {MakeField("vec1", MetricType::L2)};
auto result =
reranker::rerank(reranker::WeightedParams{{1.0}}, results, fields,
/*topn=*/2);
ASSERT_TRUE(result.has_value());
ASSERT_EQ(result.value().size(), 2u);
}
// ==================== Callback Tests ====================
TEST(RerankCallbackTest, BasicCallback) {
// Simple callback that returns docs sorted by score descending, limited to
// topn.
reranker::CallbackParams::Callback cb =
[](const std::vector<DocPtrList> &results,
const std::vector<FieldSchema::Ptr> & /*fields*/,
int topn) -> DocPtrList {
DocPtrList all_docs;
for (const auto &docs : results) {
for (const auto &doc : docs) {
all_docs.push_back(doc);
}
}
std::sort(all_docs.begin(), all_docs.end(),
[](const Doc::Ptr &a, const Doc::Ptr &b) {
return a->score() > b->score();
});
if (static_cast<int>(all_docs.size()) > topn) {
all_docs.resize(topn);
}
return all_docs;
};
std::vector<DocPtrList> results;
results.push_back({MakeDoc("a", 0.5f), MakeDoc("b", 0.9f)});
results.push_back({MakeDoc("c", 0.7f)});
auto result =
reranker::rerank(reranker::CallbackParams{cb}, results, /*fields=*/{},
/*topn=*/10);
ASSERT_TRUE(result.has_value());
auto &out = result.value();
ASSERT_EQ(out.size(), 3u);
// Should be sorted by score descending.
EXPECT_EQ(out[0]->pk(), "b");
EXPECT_EQ(out[1]->pk(), "c");
EXPECT_EQ(out[2]->pk(), "a");
}
TEST(RerankCallbackTest, EmptyCallbackError) {
reranker::CallbackParams params; // callback is empty
std::vector<DocPtrList> results;
results.push_back({MakeDoc("a", 0.5f)});
auto result = reranker::rerank(params, results, /*fields=*/{}, /*topn=*/10);
ASSERT_FALSE(result.has_value());
}
+43
View File
@@ -0,0 +1,43 @@
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
if(APPLE)
set(APPLE_FRAMEWORK_LIBS
-framework CoreFoundation
-framework CoreGraphics
-framework CoreData
-framework CoreText
-framework Security
-framework Foundation
-Wl,-U,_MallocExtension_ReleaseFreeMemory
-Wl,-U,_ProfilerStart
-Wl,-U,_ProfilerStop
-Wl,-U,_RegisterThriftProtocol
)
endif()
file(GLOB ALL_TEST_SRCS *_test.cc)
foreach(CC_SRCS ${ALL_TEST_SRCS})
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
cc_gmock(
NAME ${CC_TARGET} STRICT
LIBS zvec_common
zvec_proto
zvec_sqlengine
zvec
zvec_ailego
core_metric
core_utility
core_quantizer
core_knn_hnsw core_knn_hnsw_sparse sparsehash
core_knn_flat core_knn_flat_sparse core_knn_ivf
core_knn_hnsw_rabitq core_mix_reducer
${CMAKE_THREAD_LIBS_INIT}
${CMAKE_DL_LIBS}
SRCS ${CC_SRCS}
INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/db ${PROJECT_ROOT_DIR}/src/db/common
LDFLAGS ${APPLE_FRAMEWORK_LIBS}
)
cc_test_suite(zvec_sqlengine ${CC_TARGET})
endforeach()
+484
View File
@@ -0,0 +1,484 @@
// 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 <cstdint>
#include <cstdlib>
#include <memory>
#include <arrow/api.h>
#include <arrow/io/api.h>
#include <arrow/ipc/api.h>
#include <gtest/gtest.h>
#include "db/common/file_helper.h"
#include "db/index/segment/segment.h"
#include "db/sqlengine/sqlengine.h"
#include "zvec/db/index_params.h"
#include "zvec/db/schema.h"
#include "zvec/db/type.h"
#include "test_helper.h"
namespace zvec::sqlengine {
static Doc create_doc(const uint64_t doc_id) {
Doc new_doc;
new_doc.set_pk("pk_" + std::to_string(doc_id));
new_doc.set_doc_id(doc_id);
auto size = doc_id % 100;
if (size > 0) {
std::vector<std::string> str_array;
std::vector<int32_t> i32_array;
std::vector<int64_t> i64_array;
std::vector<uint32_t> u32_array;
std::vector<uint64_t> u64_array;
std::vector<float> fp32_array;
std::vector<double> fp64_array;
std::vector<bool> bool_array;
for (uint32_t i = 1; i <= size; i++) {
i32_array.push_back(i);
i64_array.push_back(i);
u32_array.push_back(i);
u64_array.push_back(i);
fp32_array.push_back(i);
fp64_array.push_back(i);
bool_array.push_back(i % 2 == 0);
str_array.push_back("name" + std::to_string(i));
}
new_doc.set("i32_array", i32_array);
new_doc.set("i64_array", i64_array);
new_doc.set("u32_array", u32_array);
new_doc.set("u64_array", u64_array);
new_doc.set("fp32_array", fp32_array);
new_doc.set("fp64_array", fp64_array);
new_doc.set("bool_array", bool_array);
new_doc.set("str_array", str_array);
}
return new_doc;
}
class ContainTest : public testing::Test {
protected:
static void SetUpTestSuite() {
FileHelper::RemoveDirectory(seg_path_);
FileHelper::CreateDirectory(seg_path_);
auto invert_params = std::make_shared<InvertIndexParams>(true);
collection_schema_ = std::make_shared<CollectionSchema>(
"test_collection",
std::vector<FieldSchema::Ptr>{
std::make_shared<FieldSchema>("str_array", DataType::ARRAY_STRING,
true, nullptr),
std::make_shared<FieldSchema>("i32_array", DataType::ARRAY_INT32,
true, nullptr),
std::make_shared<FieldSchema>("i64_array", DataType::ARRAY_INT64,
true, nullptr),
std::make_shared<FieldSchema>("u32_array", DataType::ARRAY_UINT32,
true, nullptr),
std::make_shared<FieldSchema>("u64_array", DataType::ARRAY_UINT64,
true, nullptr),
std::make_shared<FieldSchema>("fp32_array", DataType::ARRAY_FLOAT,
true, nullptr),
std::make_shared<FieldSchema>("fp64_array", DataType::ARRAY_DOUBLE,
true, nullptr),
std::make_shared<FieldSchema>("bool_array", DataType::ARRAY_BOOL,
true, nullptr),
});
auto segment = create_segment(seg_path_, *collection_schema_);
if (segment == nullptr) {
LOG_ERROR("create segment failed");
EXPECT_TRUE(segment != nullptr);
std::exit(EXIT_FAILURE);
}
auto status = InsertDoc(segment, 0, 10000, &create_doc);
if (!status.ok()) {
LOG_ERROR("insert doc failed: %s", status.c_str());
EXPECT_TRUE(status.ok());
std::exit(EXIT_FAILURE);
}
segments_.push_back(segment);
}
static void TearDownTestSuite() {
segments_.clear();
FileHelper::RemoveDirectory(seg_path_);
}
protected:
static inline std::string seg_path_ = "./test_collection";
static inline CollectionSchema::Ptr collection_schema_;
static inline std::vector<Segment::Ptr> segments_;
};
TEST_F(ContainTest, ContainAllInt32) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "i32_array contain_all (";
for (int i = 1; i <= 32; i++) {
query.filter_ += std::to_string(i);
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 32) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAllInt64) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "i64_array contain_all (";
for (int i = 1; i <= 32; i++) {
query.filter_ += std::to_string(i);
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 32) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAllUint32) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "u32_array contain_all (";
for (int i = 1; i <= 32; i++) {
query.filter_ += std::to_string(i);
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 32) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAllUint64) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "u64_array contain_all (";
for (int i = 1; i <= 32; i++) {
query.filter_ += std::to_string(i);
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 32) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAllFp32) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "fp32_array contain_all (";
for (int i = 1; i <= 32; i++) {
query.filter_ += std::to_string(i);
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 32) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAllFp64) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "fp64_array contain_all (";
for (int i = 1; i <= 32; i++) {
query.filter_ += std::to_string(i);
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 32) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAllString) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "str_array contain_all (";
for (int i = 1; i <= 32; i++) {
query.filter_ += "'name" + std::to_string(i) + "'";
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 32) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAnyInt32) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "i32_array contain_any (98,99,100)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 98; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 98) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAnyInt64) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "i64_array contain_any (98,99,100)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 98; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 98) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAnyUint32) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "u32_array contain_any (98,99,100)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 98; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 98) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAnyUint64) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "u64_array contain_any (98,99,100)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 98; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 98) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAnyFp32) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "fp32_array contain_any (98,99,100)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 98; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 98) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAnyFp64) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "fp64_array contain_any (98,99,100)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 98; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 98) {
i += 1;
}
}
}
TEST_F(ContainTest, ContainAnyString) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
query.filter_ = "str_array contain_any ('name98','name99','name100')";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 98; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
i += 1;
while (i % 100 < 98) {
i += 1;
}
}
}
} // namespace zvec::sqlengine
+938
View File
@@ -0,0 +1,938 @@
// 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 <cstdint>
#include <memory>
#include <gtest/gtest.h>
#include "db/sqlengine/sqlengine.h"
#include "zvec/db/schema.h"
#include "recall_base.h"
namespace zvec::sqlengine {
class ForwardRecallTest : public RecallTest {};
TEST_F(ForwardRecallTest, Basic) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value());
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int i = 0; i < query.topk_; i++) {
auto &doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, BasicWithDocId) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.include_doc_id_ = true;
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value());
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int i = 0; i < query.topk_; i++) {
auto &doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(doc->doc_id(), i);
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, OutputNoFields) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>{};
query.topk_ = 200;
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value());
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int i = 0; i < query.topk_; i++) {
auto &doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(doc->field_names().size(), 0);
}
}
TEST_F(ForwardRecallTest, DenseVector) {
SearchQuery query;
query.output_fields_ = {"id", "dense"};
query.topk_ = 200;
query.include_vector_ = true;
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int i = 0; i < query.topk_; i++) {
auto &doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto dense = doc->get<std::vector<float>>("dense");
ASSERT_TRUE(dense.has_value());
EXPECT_EQ(dense.value().size(), 4);
for (auto v : dense.value()) {
EXPECT_FLOAT_EQ(v, (float)i);
}
}
}
TEST_F(ForwardRecallTest, SparseVector) {
SearchQuery query;
query.output_fields_ = {"id", "sparse"};
query.topk_ = 200;
query.include_vector_ = true;
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int i = 0; i < query.topk_; i++) {
auto &doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
// EXPECT_EQ(doc->field_names().size(), 2);
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto sparse =
doc->get<std::pair<std::vector<uint32_t>, std::vector<float>>>(
"sparse");
if (i % 100 == 0) {
// set with empty vector
ASSERT_FALSE(sparse.has_value());
continue;
}
ASSERT_TRUE(sparse.has_value());
const auto &[indices, values] = sparse.value();
EXPECT_EQ(indices.size(), i % 100);
EXPECT_EQ(values.size(), i % 100);
for (int j = 0; j < i % 100; j++) {
EXPECT_EQ(indices[j], j);
EXPECT_FLOAT_EQ(values[j], (float)i);
}
}
}
TEST_F(ForwardRecallTest, MultiSegment) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>();
query.topk_ = 200;
query.include_vector_ = true;
auto engine = SQLEngine::create(std::make_shared<Profiler>());
std::vector<Segment::Ptr> segments = segments_;
segments.push_back(segments_[0]);
auto ret = engine->execute(collection_schema_, query, segments);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int i = 0; i < query.topk_; i++) {
auto &doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto dense = doc->get<std::vector<float>>("dense");
ASSERT_TRUE(dense.has_value());
EXPECT_EQ(dense.value().size(), 4);
for (auto v : dense.value()) {
EXPECT_FLOAT_EQ(v, (float)i);
}
auto sparse =
doc->get<std::pair<std::vector<uint32_t>, std::vector<float>>>(
"sparse");
if (i % 100 == 0) {
// set with empty vector
ASSERT_FALSE(sparse.has_value());
continue;
}
ASSERT_TRUE(sparse.has_value());
const auto &[indices, values] = sparse.value();
EXPECT_EQ(indices.size(), i % 100);
EXPECT_EQ(values.size(), i % 100);
for (int j = 0; j < i % 100; j++) {
EXPECT_EQ(indices[j], j);
EXPECT_FLOAT_EQ(values[j], (float)i);
}
}
}
TEST_F(ForwardRecallTest, Eq) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "age = 1";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 100);
for (int j = 0, i = 1; j < (int)docs.size(); j++, i += 100) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, Gt) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "id > 1000";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int j = 0; j < query.topk_; j++) {
auto &doc = docs[j];
auto i = j + 1001;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, Ge) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "id >= 1000";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int j = 0; j < query.topk_; j++) {
auto &doc = docs[j];
auto i = j + 1000;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, Lt) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "id < 100";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
ASSERT_EQ(docs.size(), 100);
for (int j = 0, i = 0; j < (int)docs.size(); j++, i += 1) {
auto &doc = docs[j];
EXPECT_EQ(i, doc->get<uint64_t>("id"));
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, Le) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "id <= 100";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
ASSERT_EQ(docs.size(), 101);
for (int j = 0, i = 0; j < (int)docs.size(); j++, i += 1) {
auto &doc = docs[j];
EXPECT_EQ(i, doc->get<uint64_t>("id"));
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, And) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "id <= 100 and id > 50";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
ASSERT_EQ(docs.size(), 50);
for (int j = 0, i = 51; j < (int)docs.size(); j++, i += 1) {
auto &doc = docs[j];
EXPECT_EQ(i, doc->get<uint64_t>("id"));
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, Or) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "id < 100 or id > 200";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
ASSERT_EQ(docs.size(), 200);
for (int j = 0; j < (int)docs.size(); j++) {
int i = j < 100 ? j : j + 101;
auto &doc = docs[j];
EXPECT_EQ(i, doc->get<uint64_t>("id"));
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, StrEq) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "name = 'user_1'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 100);
for (int j = 0, i = 1; j < (int)docs.size(); j++, i += 100) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, StrGe) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "name >= 'user_1'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 0; j < (int)docs.size(); j++, i += 1) {
if (i % 100 == 0) {
i += 1;
}
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, StrIn) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "name IN ('user_1', 'user_2')";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 1; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
if (i % 100 == 1) {
i += 1;
} else if (i % 100 == 2) {
i += 99;
}
}
}
TEST_F(ForwardRecallTest, StrNotIn) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "name NOT IN ('user_1', 'user_2')";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 0; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
if (i % 100 == 0) {
i += 3;
} else {
i += 1;
}
}
}
TEST_F(ForwardRecallTest, StrLike) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "name like 'user_9%'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 9; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
if (i % 100 == 9) {
i += 81;
} else if (i % 100 == 99) {
i += 10;
} else {
i += 1;
}
}
}
TEST_F(ForwardRecallTest, IsNull) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "optional_age is null";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 100);
for (int j = 0, i = 0; j < (int)docs.size(); j++, i += 100) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, IsNotNull) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "optional_age is not null";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 0; j < (int)docs.size(); j++, i += 1) {
if (i % 100 == 0) {
i += 1;
}
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, IsNullNoResult) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "age is null";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
ASSERT_EQ(docs.size(), 0);
}
TEST_F(ForwardRecallTest, ContainAll) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "category_set contain_all (";
for (int i = 1; i <= 32; i++) {
query.filter_ += std::to_string(i);
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
while (i % 100 < 32) {
i += 1;
}
}
}
TEST_F(ForwardRecallTest, NotContainAll) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "category_set not contain_all (";
for (int i = 1; i <= 32; i++) {
query.filter_ += std::to_string(i);
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 1; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
// i % 100 == 0 has null category
while (i % 100 >= 32 || i % 100 == 0) {
i += 1;
}
}
}
TEST_F(ForwardRecallTest, ContainAny) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "category_set contain_any (98,99,100)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 98; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
while (i % 100 < 98) {
i += 1;
}
}
}
TEST_F(ForwardRecallTest, NotContainAny) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "category_set not contain_any (98,99,100)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 1; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
// i % 100 == 0 has null category
while (i % 100 >= 98 || i % 100 == 0) {
i += 1;
}
}
}
TEST_F(ForwardRecallTest, BoolContainAll) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "bool_array contain_all (true, false)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 0; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 3;
}
}
TEST_F(ForwardRecallTest, BoolContainAny) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "bool_array contain_any (true)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 0; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
if (i % 3 == 2) {
i += 1;
}
}
}
TEST_F(ForwardRecallTest, ContainAllEmptySet) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "category_set contain_all ()";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 1; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
while (i % 100 == 0) {
i += 1;
}
}
}
TEST_F(ForwardRecallTest, NotContainAllEmptySet) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "category_set not contain_all ()";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 0);
}
TEST_F(ForwardRecallTest, ContainAnyEmptySet) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "category_set contain_any ()";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 0);
}
TEST_F(ForwardRecallTest, NotContainAnyEmptySet) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "category_set not contain_any ()";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 1; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
while (i % 100 == 0) {
i += 1;
}
}
}
TEST_F(ForwardRecallTest, BoolEqTrue) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "bool = TRuE";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 100);
for (int j = 0, i = 0; j < (int)docs.size(); j++, i += 100) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(ForwardRecallTest, BoolEqFalse) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "bool = false";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 1; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
if (i % 100 == 0) {
i += 1;
}
}
}
TEST_F(ForwardRecallTest, ArrayLengthEq) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "array_length(category_set) = 32";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 100);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 100;
}
}
TEST_F(ForwardRecallTest, ArrayLengthGe) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "array_length(category_set) >= 32";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
while (i % 100 < 32) {
i += 1;
}
}
}
} // namespace zvec::sqlengine
@@ -0,0 +1,233 @@
// 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 <memory>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "db/common/file_helper.h"
#include "db/index/common/version_manager.h"
#include "db/index/segment/segment.h"
#include "db/sqlengine/sqlengine.h"
#include "zvec/db/doc.h"
#include "zvec/db/index_params.h"
#include "zvec/db/schema.h"
#include "zvec/db/type.h"
namespace zvec::sqlengine {
// Multi-segment FTS recall regression:
//
// The planner's SegmentNode drains per-segment readers in LIFO order, so the
// per-segment BM25 ordering is *not* preserved across the merged stream. The
// planner must therefore add a global order_by on the score column for FTS,
// mirroring what it already does for vector queries.
//
// To make the regression observable we engineer the two segments so that
// * segments_[0] (read LAST) holds the globally highest-scoring doc, and
// * segments_[1] (read FIRST) holds many low-scoring docs.
//
// Per-segment BM25 stats (rare term -> high IDF in segments_[0], common term
// -> low IDF in segments_[1]) guarantee s0_0 outranks every doc in
// segments_[1]. Without the global sort the first doc in the merged stream is
// the much lower-scoring s1_*, which breaks both the descending invariant and
// topk truncation.
class FtsMultiSegmentTest : public ::testing::Test {
protected:
static void SetUpTestSuite() {
FileHelper::RemoveDirectory(root_path_);
FileHelper::CreateDirectory(root_path_);
build_schema();
// segments_[0]: only one doc contains "apple" but with very high TF and
// very low df (rare term) -> high BM25.
auto seg0 = create_segment(root_path_ + "/seg0", "fts_ms_seg0");
ASSERT_NE(seg0, nullptr);
insert_docs(seg0, /*pk_prefix=*/"s0_",
{
{"apple apple apple apple apple"}, // doc 0: TF=5, df=1
{"banana"},
{"cherry"},
{"date"},
{"elderberry"},
});
// segments_[1]: all docs contain "apple" (df=N) -> very low IDF -> low
// BM25 across the board.
auto seg1 = create_segment(root_path_ + "/seg1", "fts_ms_seg1");
ASSERT_NE(seg1, nullptr);
insert_docs(seg1, /*pk_prefix=*/"s1_",
{
{"apple banana"},
{"apple cherry"},
{"apple date"},
{"apple elderberry"},
});
segments_.push_back(seg0);
segments_.push_back(seg1);
engine_ = SQLEngine::create(std::make_shared<Profiler>());
}
static void TearDownTestSuite() {
segments_.clear();
engine_.reset();
schema_.reset();
FileHelper::RemoveDirectory(root_path_);
}
Result<DocPtrList> fts_search(const std::string &query_string,
int topk = 10) {
SearchQuery vq;
vq.topk_ = topk;
vq.target_.field_name_ = "content";
FtsClause fts;
fts.query_string_ = query_string;
vq.target_.clause_ = fts;
return engine_->execute(schema_, vq, segments_);
}
private:
static void build_schema() {
auto fts_params = std::make_shared<FtsIndexParams>(
"whitespace", std::vector<std::string>{"lowercase"}, "");
schema_ = std::make_shared<CollectionSchema>(
"fts_multi_segment_test",
std::vector<FieldSchema::Ptr>{
std::make_shared<FieldSchema>("content", DataType::STRING, false,
fts_params),
// Dummy vector field keeps the schema parity with the single-
// segment FTS fixture so the analyzer/planner paths behave the
// same.
std::make_shared<FieldSchema>(
"vec", DataType::VECTOR_FP32, 4, false,
std::make_shared<FlatIndexParams>(MetricType::L2)),
});
}
static Segment::Ptr create_segment(const std::string &seg_path,
const std::string &name) {
FileHelper::CreateDirectory(seg_path);
auto segment_meta = std::make_shared<SegmentMeta>();
segment_meta->set_id(0);
auto id_map = IDMap::CreateAndOpen(name, seg_path + "/id_map", true, false);
auto delete_store = std::make_shared<DeleteStore>(name);
Version v1;
v1.set_schema(*schema_);
std::string v_path = seg_path + "/manifest";
FileHelper::CreateDirectory(v_path);
auto vm = VersionManager::Create(v_path, v1);
if (!vm.has_value()) {
return nullptr;
}
BlockMeta mem_block;
mem_block.id_ = 0;
mem_block.type_ = BlockType::SCALAR;
mem_block.min_doc_id_ = 0;
mem_block.max_doc_id_ = 0;
mem_block.doc_count_ = 0;
segment_meta->set_writing_forward_block(mem_block);
SegmentOptions options;
options.read_only_ = false;
options.enable_mmap_ = true;
options.max_buffer_size_ = 256 * 1024;
auto result = Segment::CreateAndOpen(seg_path, *schema_, 0, 0, id_map,
delete_store, vm.value(), options);
if (!result) {
return nullptr;
}
return result.value();
}
struct Entry {
std::string content;
};
static void insert_docs(const Segment::Ptr &segment,
const std::string &pk_prefix,
const std::vector<Entry> &entries) {
for (size_t i = 0; i < entries.size(); ++i) {
Doc doc;
doc.set_pk(pk_prefix + std::to_string(i));
doc.set_doc_id(i);
doc.set<std::string>("content", entries[i].content);
auto status = segment->Insert(doc);
ASSERT_TRUE(status.ok())
<< pk_prefix << i << " insert failed: " << status.c_str();
}
}
protected:
static inline std::string root_path_ = "./fts_multi_segment_test_collection";
static inline CollectionSchema::Ptr schema_;
static inline std::vector<Segment::Ptr> segments_;
static inline SQLEngine::Ptr engine_;
};
// The merged stream from all segments must be strictly non-increasing in
// score. Without the global order_by, segments_[1]'s low-scoring docs would
// appear before segments_[0]'s much higher-scoring s0_0, violating BM25 rank.
TEST_F(FtsMultiSegmentTest, ScoreDescendingAcrossSegments) {
auto result = fts_search("apple");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
// s0_0 + s1_0..s1_3 = 5 matches.
ASSERT_EQ(result->size(), 5u);
// s0_0 (TF=5, rare term in seg0) dominates the 4 low-IDF s1_* docs.
EXPECT_EQ((*result)[0]->pk(), "s0_0");
EXPECT_GT((*result)[0]->score(), (*result)[1]->score());
for (size_t i = 0; i + 1 < result->size(); ++i) {
EXPECT_GE((*result)[i]->score(), (*result)[i + 1]->score())
<< "score not descending at rank " << i << ": " << (*result)[i]->pk()
<< "=" << (*result)[i]->score() << " vs " << (*result)[i + 1]->pk()
<< "=" << (*result)[i + 1]->score();
}
}
// topk must cut against the globally-sorted stream. Without the fix the
// first batch surfaced from SegmentNode comes from segments_[1] (LIFO read),
// so topk=1 would silently drop the highest-scoring s0_0.
TEST_F(FtsMultiSegmentTest, TopkPicksGloballyHighestScore) {
auto result = fts_search("apple", /*topk=*/1);
ASSERT_TRUE(result.has_value()) << result.error().c_str();
ASSERT_EQ(result->size(), 1u);
EXPECT_EQ((*result)[0]->pk(), "s0_0");
}
// Sanity: a cross-segment OR query still returns the union of matches and
// stays descending across the segment boundary.
TEST_F(FtsMultiSegmentTest, CrossSegmentUnionDescending) {
// apple: 5 docs (s0_0, s1_0..s1_3). banana: s0_1 (seg0), s1_0 (seg1).
// OR-union: {s0_0, s0_1, s1_0, s1_1, s1_2, s1_3} = 6 docs.
auto result = fts_search("apple banana");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
ASSERT_EQ(result->size(), 6u);
for (size_t i = 0; i + 1 < result->size(); ++i) {
EXPECT_GE((*result)[i]->score(), (*result)[i + 1]->score())
<< "score not descending at rank " << i;
}
}
} // namespace zvec::sqlengine
+889
View File
@@ -0,0 +1,889 @@
// 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 <gtest/gtest.h>
#include "db/index/column/fts_column/fts_query_ast.h"
#include "db/index/column/fts_column/fts_types.h"
#include "db/index/column/fts_column/parser/fts_query_parser.h"
#include "db/index/column/fts_column/tokenizer/tokenizer_factory.h"
namespace zvec::fts {
// ============================================================
// Test fixture
// ============================================================
class FtsParserTest : public ::testing::Test {
protected:
void SetUp() override {
// Standard tokenizer + lowercase filter. These parser tests cover
// punctuation that standard still treats as delimiters, while CJK tests
// exercise the per-character tokens standard produces for ideographs.
FtsIndexParams params;
params.tokenizer_name = "standard";
params.filters = {"lowercase"};
pipeline_ = TokenizerFactory::create(params);
ASSERT_NE(pipeline_, nullptr);
}
FtsAstNodePtr parse(const std::string &query) {
return parser_.parse(query, pipeline_);
}
// Overload for tests that need to specify the default operator explicitly.
FtsAstNodePtr parse(const std::string &query, FtsDefaultOperator default_op) {
return parser_.parse(query, pipeline_, default_op);
}
const std::string &err_msg() {
return parser_.err_msg();
}
// Helpers for type-safe downcasting
static const TermNode &as_term(const FtsAstNode &node) {
EXPECT_EQ(node.type(), FtsNodeType::TERM);
return static_cast<const TermNode &>(node);
}
static const PhraseNode &as_phrase(const FtsAstNode &node) {
EXPECT_EQ(node.type(), FtsNodeType::PHRASE);
return static_cast<const PhraseNode &>(node);
}
static const AndNode &as_and(const FtsAstNode &node) {
EXPECT_EQ(node.type(), FtsNodeType::AND);
return static_cast<const AndNode &>(node);
}
static const OrNode &as_or(const FtsAstNode &node) {
EXPECT_EQ(node.type(), FtsNodeType::OR);
return static_cast<const OrNode &>(node);
}
private:
FtsQueryParser parser_;
TokenizerPipelinePtr pipeline_;
};
// ============================================================
// Single term
// ============================================================
TEST_F(FtsParserTest, SingleTerm) {
auto ast = parse("vector");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
const auto &term = as_term(*ast);
EXPECT_EQ(term.term, "vector");
EXPECT_FALSE(term.must);
EXPECT_FALSE(term.must_not);
}
TEST_F(FtsParserTest, SingleTermNumeric) {
auto ast = parse("2024");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_EQ(as_term(*ast).term, "2024");
}
TEST_F(FtsParserTest, SingleTermWithHyphen) {
// The lexer's REGULAR_ID rule keeps hyphenated text as one token, but the
// standard tokenizer on the parser side splits this hyphen delimiter. With
// the default OR operator the term decomposes into Or[full, text] so query
// segmentation matches the index segmentation.
auto ast = parse("full-text");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 2u);
EXPECT_EQ(as_term(*or_node.children[0]).term, "full");
EXPECT_EQ(as_term(*or_node.children[1]).term, "text");
}
TEST_F(FtsParserTest, BareColonQueryIsFieldPrefixSyntax) {
auto ast = parse("host:port");
EXPECT_EQ(ast, nullptr);
EXPECT_EQ(err_msg(), "field-prefixed queries are not supported");
}
// ============================================================
// Must (+) and must_not (-/NOT) modifiers
// ============================================================
TEST_F(FtsParserTest, MustModifier) {
auto ast = parse("+vector");
ASSERT_NE(ast, nullptr);
const auto &term = as_term(*ast);
EXPECT_EQ(term.term, "vector");
EXPECT_TRUE(term.must);
EXPECT_FALSE(term.must_not);
}
TEST_F(FtsParserTest, MustNotModifierMinus) {
// "-slow" is lexed as a single REGULAR_ID token (hyphen is part of the id).
// To express must_not, use a space: "- slow" -> MINUS_SIGN + REGULAR_ID.
auto ast = parse("- slow");
ASSERT_NE(ast, nullptr);
const auto &term = as_term(*ast);
EXPECT_EQ(term.term, "slow");
EXPECT_FALSE(term.must);
EXPECT_TRUE(term.must_not);
}
TEST_F(FtsParserTest, MustNotModifierMinusNoSpace) {
// "-slow" without space: FtsLexer treats '-' as MINUS_SIGN modifier,
// so "-slow" is parsed as must_not:slow (same as "- slow").
auto ast = parse("-slow");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_EQ(as_term(*ast).term, "slow");
EXPECT_TRUE(as_term(*ast).must_not);
}
TEST_F(FtsParserTest, MustNotModifierNot) {
// NOT is now a strict binary operator (`a NOT b` <=> `a AND NOT b`).
// A leading `NOT a` is therefore a syntax error — there is no left-hand
// operand for NOT to subtract from.
auto ast = parse("NOT slow");
EXPECT_EQ(ast, nullptr);
EXPECT_FALSE(err_msg().empty());
}
// ============================================================
// Phrase query
// ============================================================
TEST_F(FtsParserTest, DoubleQuotedPhrase) {
auto ast = parse("\"exact phrase\"");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::PHRASE);
const auto &phrase = as_phrase(*ast);
ASSERT_EQ(phrase.terms.size(), 2u);
EXPECT_EQ(phrase.terms[0], "exact");
EXPECT_EQ(phrase.terms[1], "phrase");
EXPECT_FALSE(phrase.must);
EXPECT_FALSE(phrase.must_not);
}
TEST_F(FtsParserTest, SingleQuotedPhrase) {
// Single-quoted strings are not supported as phrase queries (no SQUOTA_STRING
// token). The lexer's TERM rule absorbs "'hello", "world", and "'" as
// individual term tokens, so the query parses as an implicit OR of terms.
auto ast = parse("'hello world'");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
}
TEST_F(FtsParserTest, PhraseWithMustModifier) {
auto ast = parse("+\"exact phrase\"");
ASSERT_NE(ast, nullptr);
const auto &phrase = as_phrase(*ast);
EXPECT_TRUE(phrase.must);
EXPECT_FALSE(phrase.must_not);
}
TEST_F(FtsParserTest, PhraseWithMustNotModifier) {
auto ast = parse("-\"bad phrase\"");
ASSERT_NE(ast, nullptr);
const auto &phrase = as_phrase(*ast);
EXPECT_FALSE(phrase.must);
EXPECT_TRUE(phrase.must_not);
}
TEST_F(FtsParserTest, PhraseWithThreeWords) {
auto ast = parse("\"one two three\"");
ASSERT_NE(ast, nullptr);
const auto &phrase = as_phrase(*ast);
ASSERT_EQ(phrase.terms.size(), 3u);
EXPECT_EQ(phrase.terms[0], "one");
EXPECT_EQ(phrase.terms[1], "two");
EXPECT_EQ(phrase.terms[2], "three");
}
// ============================================================
// Explicit OR
// ============================================================
TEST_F(FtsParserTest, ExplicitOr) {
auto ast = parse("cat OR dog");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 2u);
EXPECT_EQ(as_term(*or_node.children[0]).term, "cat");
EXPECT_EQ(as_term(*or_node.children[1]).term, "dog");
}
TEST_F(FtsParserTest, MultipleOr) {
auto ast = parse("a OR b OR c");
ASSERT_NE(ast, nullptr);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 3u);
}
// ============================================================
// Explicit AND
// ============================================================
TEST_F(FtsParserTest, ExplicitAnd) {
auto ast = parse("cat AND dog");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 2u);
EXPECT_EQ(as_term(*and_node.children[0]).term, "cat");
EXPECT_EQ(as_term(*and_node.children[1]).term, "dog");
}
TEST_F(FtsParserTest, MultipleAnd) {
auto ast = parse("a AND b AND c");
ASSERT_NE(ast, nullptr);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 3u);
}
// ============================================================
// Operator precedence: AND binds tighter than OR
// ============================================================
TEST_F(FtsParserTest, AndBindsTighterThanOr) {
// "a OR b AND c" should parse as "a OR (b AND c)"
auto ast = parse("a OR b AND c");
ASSERT_NE(ast, nullptr);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 2u);
// Left child: term "a"
EXPECT_EQ(as_term(*or_node.children[0]).term, "a");
// Right child: AND(b, c)
const auto &and_node = as_and(*or_node.children[1]);
ASSERT_EQ(and_node.children.size(), 2u);
EXPECT_EQ(as_term(*and_node.children[0]).term, "b");
EXPECT_EQ(as_term(*and_node.children[1]).term, "c");
}
// ============================================================
// Implicit adjacency (seqExpr / default operator)
// ============================================================
TEST_F(FtsParserTest, ImplicitAdjacency) {
// Adjacent terms without explicit operator: "a b" -> seqExpr -> OR(a, b)
auto ast = parse("a b");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 2u);
EXPECT_EQ(as_term(*or_node.children[0]).term, "a");
EXPECT_EQ(as_term(*or_node.children[1]).term, "b");
}
TEST_F(FtsParserTest, ImplicitAdjacencyThreeTerms) {
auto ast = parse("a b c");
ASSERT_NE(ast, nullptr);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 3u);
}
TEST_F(FtsParserTest, ImplicitAdjacencyWithModifiers) {
// "+a - b" -> seqExpr -> OR(must:a, must_not:b)
// Note: "-b" (no space) is lexed as a single REGULAR_ID; use "- b" for
// must_not.
auto ast = parse("+a - b");
ASSERT_NE(ast, nullptr);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 2u);
EXPECT_TRUE(as_term(*or_node.children[0]).must);
EXPECT_TRUE(as_term(*or_node.children[1]).must_not);
}
// ============================================================
// Parentheses grouping
// ============================================================
TEST_F(FtsParserTest, Parentheses) {
// "(a OR b) AND c"
auto ast = parse("(a OR b) AND c");
ASSERT_NE(ast, nullptr);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 2u);
// Left: OR(a, b)
const auto &or_node = as_or(*and_node.children[0]);
ASSERT_EQ(or_node.children.size(), 2u);
// Right: term c
EXPECT_EQ(as_term(*and_node.children[1]).term, "c");
}
TEST_F(FtsParserTest, NestedParentheses) {
auto ast = parse("((a OR b) AND c) OR d");
ASSERT_NE(ast, nullptr);
const auto &outer_or = as_or(*ast);
ASSERT_EQ(outer_or.children.size(), 2u);
EXPECT_EQ(as_term(*outer_or.children[1]).term, "d");
}
// ============================================================
// Mixed complex queries
// ============================================================
TEST_F(FtsParserTest, MixedTermAndPhrase) {
// "+vector - slow \"exact phrase\""
// Note: use "- slow" (with space) so MINUS_SIGN is a separate token.
auto ast = parse("+vector - slow \"exact phrase\"");
ASSERT_NE(ast, nullptr);
// Four adjacent items -> seqExpr -> OR(must:vector, must_not:slow, phrase)
// Actually: +vector and - slow and phrase are three unary nodes in seqExpr
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 3u);
EXPECT_TRUE(as_term(*or_node.children[0]).must);
EXPECT_EQ(as_term(*or_node.children[0]).term, "vector");
EXPECT_TRUE(as_term(*or_node.children[1]).must_not);
EXPECT_EQ(as_term(*or_node.children[1]).term, "slow");
EXPECT_EQ(or_node.children[2]->type(), FtsNodeType::PHRASE);
}
TEST_F(FtsParserTest, AndWithPhrase) {
auto ast = parse("\"machine learning\" AND model");
ASSERT_NE(ast, nullptr);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 2u);
EXPECT_EQ(and_node.children[0]->type(), FtsNodeType::PHRASE);
EXPECT_EQ(as_term(*and_node.children[1]).term, "model");
}
TEST_F(FtsParserTest, ComplexBooleanQuery) {
// "a AND b OR c AND d" -> (a AND b) OR (c AND d)
auto ast = parse("a AND b OR c AND d");
ASSERT_NE(ast, nullptr);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 2u);
const auto &left_and = as_and(*or_node.children[0]);
ASSERT_EQ(left_and.children.size(), 2u);
const auto &right_and = as_and(*or_node.children[1]);
ASSERT_EQ(right_and.children.size(), 2u);
}
// ============================================================
// Single-child simplification (no unnecessary wrapping)
// ============================================================
TEST_F(FtsParserTest, SingleChildNotWrapped) {
// A single term should not be wrapped in an AndNode/OrNode
auto ast = parse("hello");
ASSERT_NE(ast, nullptr);
EXPECT_EQ(ast->type(), FtsNodeType::TERM);
}
TEST_F(FtsParserTest, SinglePhraseNotWrapped) {
auto ast = parse("\"hello world\"");
ASSERT_NE(ast, nullptr);
EXPECT_EQ(ast->type(), FtsNodeType::PHRASE);
}
// ============================================================
// Error cases
// ============================================================
TEST_F(FtsParserTest, EmptyQueryReturnsNull) {
auto ast = parse("");
EXPECT_EQ(ast, nullptr);
}
TEST_F(FtsParserTest, OnlyParenthesesReturnsNull) {
auto ast = parse("()");
EXPECT_EQ(ast, nullptr);
}
TEST_F(FtsParserTest, UnclosedPhraseParsesAsTerm) {
// An unclosed double-quote causes the DQUOTA_STRING rule to fail. The
// remaining characters are absorbed by the TERM catch-all rule, so the
// query parses as a single term rather than returning nullptr.
auto ast = parse("\"unclosed phrase");
ASSERT_NE(ast, nullptr);
}
TEST_F(FtsParserTest, UnclosedParenReturnsNull) {
auto ast = parse("(a OR b");
EXPECT_EQ(ast, nullptr);
}
// ============================================================
// Empty-AST cases: grammar valid, analyzer drops every term → EmptyNode.
// ============================================================
TEST_F(FtsParserTest, PunctuationOnlyReturnsEmpty) {
auto ast = parse("!!!");
ASSERT_NE(ast, nullptr);
EXPECT_EQ(ast->type(), FtsNodeType::EMPTY);
EXPECT_TRUE(err_msg().empty());
}
TEST_F(FtsParserTest, MultiplePunctuationTermsReturnsEmpty) {
auto ast = parse("!!! ??? ...");
ASSERT_NE(ast, nullptr);
EXPECT_EQ(ast->type(), FtsNodeType::EMPTY);
EXPECT_TRUE(err_msg().empty());
}
// ============================================================
// NOT as a binary AND-NOT operator
// ============================================================
TEST_F(FtsParserTest, NotAsBinaryAndNot) {
// `foo NOT bar` <=> `foo AND NOT bar` -> And[foo, bar(must_not)]
auto ast = parse("foo NOT bar");
ASSERT_NE(ast, nullptr);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 2u);
EXPECT_EQ(as_term(*and_node.children[0]).term, "foo");
EXPECT_FALSE(and_node.children[0]->must_not);
EXPECT_EQ(as_term(*and_node.children[1]).term, "bar");
EXPECT_TRUE(and_node.children[1]->must_not);
}
TEST_F(FtsParserTest, AndAndNot) {
// `a AND NOT b` -> And[a, b(must_not)]
auto ast = parse("a AND NOT b");
ASSERT_NE(ast, nullptr);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 2u);
EXPECT_EQ(as_term(*and_node.children[0]).term, "a");
EXPECT_FALSE(and_node.children[0]->must_not);
EXPECT_EQ(as_term(*and_node.children[1]).term, "b");
EXPECT_TRUE(and_node.children[1]->must_not);
}
TEST_F(FtsParserTest, OrThenNot) {
// Precedence check: NOT shares AND's precedence (higher than OR).
// `a OR b NOT c` -> Or[a, And[b, c(must_not)]]
auto ast = parse("a OR b NOT c");
ASSERT_NE(ast, nullptr);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 2u);
EXPECT_EQ(as_term(*or_node.children[0]).term, "a");
const auto &right_and = as_and(*or_node.children[1]);
ASSERT_EQ(right_and.children.size(), 2u);
EXPECT_EQ(as_term(*right_and.children[0]).term, "b");
EXPECT_FALSE(right_and.children[0]->must_not);
EXPECT_EQ(as_term(*right_and.children[1]).term, "c");
EXPECT_TRUE(right_and.children[1]->must_not);
}
TEST_F(FtsParserTest, NotWithGroup) {
// `a NOT (b OR c)` -> And[a, Or[b, c](must_not)]
auto ast = parse("a NOT (b OR c)");
ASSERT_NE(ast, nullptr);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 2u);
EXPECT_EQ(as_term(*and_node.children[0]).term, "a");
EXPECT_FALSE(and_node.children[0]->must_not);
ASSERT_EQ(and_node.children[1]->type(), FtsNodeType::OR);
EXPECT_TRUE(and_node.children[1]->must_not);
const auto &grouped_or = as_or(*and_node.children[1]);
ASSERT_EQ(grouped_or.children.size(), 2u);
EXPECT_EQ(as_term(*grouped_or.children[0]).term, "b");
EXPECT_EQ(as_term(*grouped_or.children[1]).term, "c");
}
TEST_F(FtsParserTest, LeadingNotIsError) {
// Leading NOT has no left-hand operand and must fail to parse.
auto ast = parse("NOT a");
EXPECT_EQ(ast, nullptr);
EXPECT_FALSE(err_msg().empty());
}
TEST_F(FtsParserTest, MultipleNotsAndAnds) {
// `a AND b NOT c AND d NOT e` -> And[a, b, c(must_not), d, e(must_not)]
auto ast = parse("a AND b NOT c AND d NOT e");
ASSERT_NE(ast, nullptr);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 5u);
EXPECT_EQ(as_term(*and_node.children[0]).term, "a");
EXPECT_FALSE(and_node.children[0]->must_not);
EXPECT_EQ(as_term(*and_node.children[1]).term, "b");
EXPECT_FALSE(and_node.children[1]->must_not);
EXPECT_EQ(as_term(*and_node.children[2]).term, "c");
EXPECT_TRUE(and_node.children[2]->must_not);
EXPECT_EQ(as_term(*and_node.children[3]).term, "d");
EXPECT_FALSE(and_node.children[3]->must_not);
EXPECT_EQ(as_term(*and_node.children[4]).term, "e");
EXPECT_TRUE(and_node.children[4]->must_not);
}
// ============================================================
// +/- modifiers on parenthesised sub-expressions
// ============================================================
TEST_F(FtsParserTest, MustOnGroup) {
// `+(a OR b)` -> Or[a, b]{must=true}
auto ast = parse("+(a OR b)");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
EXPECT_TRUE(ast->must);
EXPECT_FALSE(ast->must_not);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 2u);
EXPECT_EQ(as_term(*or_node.children[0]).term, "a");
EXPECT_EQ(as_term(*or_node.children[1]).term, "b");
}
TEST_F(FtsParserTest, MustNotOnGroup) {
// `-(a AND b)` -> And[a, b]{must_not=true}
auto ast = parse("-(a AND b)");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
EXPECT_FALSE(ast->must);
EXPECT_TRUE(ast->must_not);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 2u);
EXPECT_EQ(as_term(*and_node.children[0]).term, "a");
EXPECT_EQ(as_term(*and_node.children[1]).term, "b");
}
TEST_F(FtsParserTest, MustGroupAndOther) {
// `+(a OR b) c` -> implicit-OR collapses three siblings into a single
// OrNode: Or[Or[a, b]{must=true}, c]
// (the inner OR keeps its must flag; implicit adjacency is still OR.)
auto ast = parse("+(a OR b) c");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
const auto &outer_or = as_or(*ast);
ASSERT_EQ(outer_or.children.size(), 2u);
ASSERT_EQ(outer_or.children[0]->type(), FtsNodeType::OR);
EXPECT_TRUE(outer_or.children[0]->must);
const auto &inner_or = as_or(*outer_or.children[0]);
ASSERT_EQ(inner_or.children.size(), 2u);
EXPECT_EQ(as_term(*inner_or.children[0]).term, "a");
EXPECT_EQ(as_term(*inner_or.children[1]).term, "b");
EXPECT_EQ(as_term(*outer_or.children[1]).term, "c");
}
TEST_F(FtsParserTest, NestedGroupModifier) {
// `+((a AND b) OR c)` -> the must flag attaches to the outermost OrNode.
auto ast = parse("+((a AND b) OR c)");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
EXPECT_TRUE(ast->must);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 2u);
ASSERT_EQ(or_node.children[0]->type(), FtsNodeType::AND);
EXPECT_FALSE(or_node.children[0]->must); // inner AND not affected
const auto &inner_and = as_and(*or_node.children[0]);
ASSERT_EQ(inner_and.children.size(), 2u);
EXPECT_EQ(as_term(*inner_and.children[0]).term, "a");
EXPECT_EQ(as_term(*inner_and.children[1]).term, "b");
EXPECT_EQ(as_term(*or_node.children[1]).term, "c");
}
// ============================================================
// Default operator (FtsDefaultOperator::OR / AND)
// Only adjacent bare terms (no explicit operator) are affected; explicit
// AND / OR / + / - usages keep their original semantics.
// ============================================================
TEST_F(FtsParserTest, DefaultOperatorOr_AdjacentBareTerms) {
// Backward-compat: omitting default_op or passing OR yields the original
// implicit-OR behaviour for adjacent bare terms.
auto ast = parse("vector database", FtsDefaultOperator::OR);
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 2u);
EXPECT_EQ(as_term(*or_node.children[0]).term, "vector");
EXPECT_EQ(as_term(*or_node.children[1]).term, "database");
}
TEST_F(FtsParserTest, DefaultOperatorAnd_AdjacentBareTerms) {
// With AND default, two adjacent bare terms collapse into an AndNode.
auto ast = parse("vector database", FtsDefaultOperator::AND);
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 2u);
EXPECT_EQ(as_term(*and_node.children[0]).term, "vector");
EXPECT_EQ(as_term(*and_node.children[1]).term, "database");
}
TEST_F(FtsParserTest, DefaultOperatorAnd_SingleTermUnchanged) {
// A single term should not be wrapped in an AndNode.
auto ast = parse("vector", FtsDefaultOperator::AND);
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_EQ(as_term(*ast).term, "vector");
}
TEST_F(FtsParserTest, DefaultOperatorAnd_PropagatesIntoParens) {
// Parenthesised sub-expressions inherit the same default operator.
// `(a b) c` with AND default -> And[And[a, b], c].
auto ast = parse("(a b) c", FtsDefaultOperator::AND);
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
const auto &outer_and = as_and(*ast);
ASSERT_EQ(outer_and.children.size(), 2u);
ASSERT_EQ(outer_and.children[0]->type(), FtsNodeType::AND);
const auto &inner_and = as_and(*outer_and.children[0]);
ASSERT_EQ(inner_and.children.size(), 2u);
EXPECT_EQ(as_term(*inner_and.children[0]).term, "a");
EXPECT_EQ(as_term(*inner_and.children[1]).term, "b");
EXPECT_EQ(as_term(*outer_and.children[1]).term, "c");
}
TEST_F(FtsParserTest, DefaultOperatorAnd_DoesNotOverrideExplicitOr) {
// Explicit OR has higher-level structure; default_op only changes the
// implicit adjacency inside each seqExpr.
// `a OR b c` with AND default -> Or[a, And[b, c]].
auto ast = parse("a OR b c", FtsDefaultOperator::AND);
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 2u);
EXPECT_EQ(as_term(*or_node.children[0]).term, "a");
ASSERT_EQ(or_node.children[1]->type(), FtsNodeType::AND);
const auto &inner_and = as_and(*or_node.children[1]);
ASSERT_EQ(inner_and.children.size(), 2u);
EXPECT_EQ(as_term(*inner_and.children[0]).term, "b");
EXPECT_EQ(as_term(*inner_and.children[1]).term, "c");
}
TEST_F(FtsParserTest, DefaultOperatorOr_DoesNotOverrideExplicitAnd) {
// Grammar: andExpr = seqExpr ((AND|NOT) seqExpr)*
// `a AND b c` parses as seqExpr("a") AND seqExpr("b c").
// With OR default, seqExpr("b c") -> Or[b, c].
// Result: And[a, Or[b, c]].
auto ast = parse("a AND b c", FtsDefaultOperator::OR);
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 2u);
EXPECT_EQ(as_term(*and_node.children[0]).term, "a");
ASSERT_EQ(and_node.children[1]->type(), FtsNodeType::OR);
const auto &inner_or = as_or(*and_node.children[1]);
ASSERT_EQ(inner_or.children.size(), 2u);
EXPECT_EQ(as_term(*inner_or.children[0]).term, "b");
EXPECT_EQ(as_term(*inner_or.children[1]).term, "c");
}
TEST_F(FtsParserTest, DefaultOperatorAnd_PreservesPlusMinusModifiers) {
// `+a b -c` with AND default -> And[a{must}, b, c{must_not}].
// Modifiers on individual terms are independent of default_op.
auto ast = parse("+a b -c", FtsDefaultOperator::AND);
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 3u);
const auto &t0 = as_term(*and_node.children[0]);
EXPECT_EQ(t0.term, "a");
EXPECT_TRUE(t0.must);
EXPECT_FALSE(t0.must_not);
const auto &t1 = as_term(*and_node.children[1]);
EXPECT_EQ(t1.term, "b");
EXPECT_FALSE(t1.must);
EXPECT_FALSE(t1.must_not);
const auto &t2 = as_term(*and_node.children[2]);
EXPECT_EQ(t2.term, "c");
EXPECT_FALSE(t2.must);
EXPECT_TRUE(t2.must_not);
}
// ============================================================
// Pipeline-aware tokenization (phrase / bare term split through pipeline)
// ============================================================
TEST_F(FtsParserTest, MultiTokenBareTermAndDefaultGroupsAsAnd) {
// `full-text` lexes as one REGULAR_ID, but standard splits it into
// ["full", "text"]. With AND default operator the two tokens combine into
// an AndNode rather than the OR returned by the OR-default test above.
auto ast = parse("full-text", FtsDefaultOperator::AND);
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::AND);
const auto &and_node = as_and(*ast);
ASSERT_EQ(and_node.children.size(), 2u);
EXPECT_EQ(as_term(*and_node.children[0]).term, "full");
EXPECT_EQ(as_term(*and_node.children[1]).term, "text");
}
TEST_F(FtsParserTest, MultiTokenBareTermPreservesMustModifier) {
// `+full-text` -> Or[full, text] with must=true on the composite root.
auto ast = parse("+full-text");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::OR);
EXPECT_TRUE(ast->must);
EXPECT_FALSE(ast->must_not);
const auto &or_node = as_or(*ast);
ASSERT_EQ(or_node.children.size(), 2u);
EXPECT_EQ(as_term(*or_node.children[0]).term, "full");
EXPECT_EQ(as_term(*or_node.children[1]).term, "text");
}
TEST_F(FtsParserTest, PhraseTokensRunThroughPipeline) {
// The phrase body is tokenized exactly like document text. With the
// standard tokenizer, comma and exclamation delimiters collapse so
// "machine, learning!" becomes ["machine", "learning"].
auto ast = parse("\"machine, learning!\"");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::PHRASE);
const auto &phrase = as_phrase(*ast);
ASSERT_EQ(phrase.terms.size(), 2u);
EXPECT_EQ(phrase.terms[0], "machine");
EXPECT_EQ(phrase.terms[1], "learning");
}
TEST_F(FtsParserTest, PhraseCanSearchLiteralColonToken) {
auto ast = parse("\"host:port\"");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::PHRASE);
const auto &phrase = as_phrase(*ast);
ASSERT_EQ(phrase.terms.size(), 1u);
EXPECT_EQ(phrase.terms[0], "host:port");
}
TEST_F(FtsParserTest, PhraseLowercaseFilterApplies) {
// The lowercase filter is part of the pipeline so phrase tokens come back
// lowercased even when the input mixed case.
auto ast = parse("\"Machine LEARNING\"");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::PHRASE);
const auto &phrase = as_phrase(*ast);
ASSERT_EQ(phrase.terms.size(), 2u);
EXPECT_EQ(phrase.terms[0], "machine");
EXPECT_EQ(phrase.terms[1], "learning");
}
TEST_F(FtsParserTest, AllPunctuationPhraseYieldsEmptyTerms) {
// Pure non-alnum content is filtered out entirely. The phrase node still
// exists but carries zero terms; the search engine treats this as
// "match nothing" without crashing.
auto ast = parse("\"!!! ???\"");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::PHRASE);
EXPECT_TRUE(as_phrase(*ast).terms.empty());
}
// ============================================================
// Unescape: backslash removal for TERM and PHRASE paths.
// Uses WhitespaceTokenizer (no filter) so that special characters are
// preserved in tokens — this lets us observe whether unescape() actually
// stripped the backslashes.
// ============================================================
class FtsParserUnescapeTest : public ::testing::Test {
protected:
void SetUp() override {
FtsIndexParams params;
params.tokenizer_name = "whitespace";
params.filters = {};
pipeline_ = TokenizerFactory::create(params);
ASSERT_NE(pipeline_, nullptr);
}
FtsAstNodePtr parse(const std::string &query) {
return parser_.parse(query, pipeline_);
}
static const TermNode &as_term(const FtsAstNode &node) {
EXPECT_EQ(node.type(), FtsNodeType::TERM);
return static_cast<const TermNode &>(node);
}
static const PhraseNode &as_phrase(const FtsAstNode &node) {
EXPECT_EQ(node.type(), FtsNodeType::PHRASE);
return static_cast<const PhraseNode &>(node);
}
private:
FtsQueryParser parser_;
TokenizerPipelinePtr pipeline_;
};
TEST_F(FtsParserUnescapeTest, TermEscapedPlusBecomesLiteralPlus) {
// Lexer token: C\+\+ (with backslashes). After unescape: C++.
// WhitespaceTokenizer preserves the '+' in the token text.
auto ast = parse(R"(C\+\+)");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_EQ(as_term(*ast).term, "C++");
}
TEST_F(FtsParserUnescapeTest, TermEscapedMinusBecomesLiteralMinus) {
// "a\-b" after unescape → "a-b" kept intact by whitespace tokenizer.
auto ast = parse(R"(a\-b)");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_EQ(as_term(*ast).term, "a-b");
}
TEST_F(FtsParserUnescapeTest, TermEscapedBackslashBecomesLiteralBackslash) {
// "path\\dir" — lexer sees ESCAPED_CHAR(\\), unescape yields "path\dir".
auto ast = parse(R"(path\\dir)");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::TERM);
EXPECT_EQ(as_term(*ast).term, "path\\dir");
}
TEST_F(FtsParserUnescapeTest, PhraseEscapedQuoteBecomesLiteralQuote) {
// Phrase: "hello \"world\"" — after strip_quotes + unescape:
// 'hello "world"' — whitespace tokenizer splits on space to:
// ["hello", "\"world\""]
auto ast = parse(R"("hello \"world\"")");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::PHRASE);
const auto &phrase = as_phrase(*ast);
ASSERT_EQ(phrase.terms.size(), 2u);
EXPECT_EQ(phrase.terms[0], "hello");
EXPECT_EQ(phrase.terms[1], "\"world\"");
}
TEST_F(FtsParserUnescapeTest, PhraseEscapedBackslashBecomesLiteral) {
// Phrase: "a\\b" — after strip+unescape: "a\b" (one backslash, no space),
// whitespace tokenizer keeps it as single token.
auto ast = parse(R"("a\\b")");
ASSERT_NE(ast, nullptr);
ASSERT_EQ(ast->type(), FtsNodeType::PHRASE);
const auto &phrase = as_phrase(*ast);
ASSERT_EQ(phrase.terms.size(), 1u);
EXPECT_EQ(phrase.terms[0], "a\\b");
}
} // namespace zvec::fts
+824
View File
@@ -0,0 +1,824 @@
// 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 <cstdint>
#include <memory>
#include <set>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "db/common/file_helper.h"
#include "db/index/common/version_manager.h"
#include "db/index/segment/segment.h"
#include "db/sqlengine/sqlengine.h"
#include "zvec/db/doc.h"
#include "zvec/db/index_params.h"
#include "zvec/db/query_params.h"
#include "zvec/db/schema.h"
#include "zvec/db/type.h"
namespace zvec::sqlengine {
// ============================================================
// FTS Recall Test fixture (real Segment + SQLEngine::execute via SearchQuery)
// ============================================================
class FtsRecallTest : public ::testing::Test {
protected:
static void SetUpTestSuite() {
FileHelper::RemoveDirectory(seg_path_);
FileHelper::CreateDirectory(seg_path_);
build_schema();
auto segment = create_segment();
ASSERT_NE(segment, nullptr);
insert_docs(segment);
segments_.push_back(segment);
engine_ = SQLEngine::create(std::make_shared<Profiler>());
}
static void TearDownTestSuite() {
segments_.clear();
engine_.reset();
schema_.reset();
FileHelper::RemoveDirectory(seg_path_);
}
// Helper: execute FTS query_string search via SearchQuery
Result<DocPtrList> fts_search(const std::string &query_string,
int topk = 10) {
SearchQuery vq;
vq.topk_ = topk;
vq.target_.field_name_ = "content";
FtsClause fts;
fts.query_string_ = query_string;
vq.target_.clause_ = fts;
return engine_->execute(schema_, vq, segments_);
}
// Helper: execute FTS match_string search via SearchQuery
Result<DocPtrList> fts_match(const std::string &match_string,
const std::string &default_op = "",
int topk = 10) {
SearchQuery vq;
vq.topk_ = topk;
vq.target_.field_name_ = "content";
FtsClause fts;
fts.match_string_ = match_string;
vq.target_.clause_ = fts;
if (!default_op.empty()) {
auto fts_qp = std::make_shared<zvec::FtsQueryParams>();
fts_qp->set_default_operator(default_op);
vq.target_.query_params_ = fts_qp;
}
return engine_->execute(schema_, vq, segments_);
}
// Helper: execute FTS query_string with default_operator via SearchQuery
Result<DocPtrList> fts_query_with_op(const std::string &query_string,
const std::string &default_op,
int topk = 10) {
SearchQuery vq;
vq.topk_ = topk;
vq.target_.field_name_ = "content";
FtsClause fts;
fts.query_string_ = query_string;
vq.target_.clause_ = fts;
auto fts_qp = std::make_shared<zvec::FtsQueryParams>();
fts_qp->set_default_operator(default_op);
vq.target_.query_params_ = fts_qp;
return engine_->execute(schema_, vq, segments_);
}
// Helper: execute FTS query_string with WHERE filter via SearchQuery
Result<DocPtrList> fts_search_with_filter(const std::string &query_string,
const std::string &filter,
int topk = 10) {
SearchQuery vq;
vq.topk_ = topk;
vq.target_.field_name_ = "content";
vq.filter_ = filter;
FtsClause fts;
fts.query_string_ = query_string;
vq.target_.clause_ = fts;
return engine_->execute(schema_, vq, segments_);
}
private:
static void build_schema() {
auto fts_params = std::make_shared<FtsIndexParams>(
"whitespace", std::vector<std::string>{"lowercase"}, "");
auto invert_params = std::make_shared<InvertIndexParams>(true);
schema_ = std::make_shared<CollectionSchema>(
"fts_recall_test",
std::vector<FieldSchema::Ptr>{
std::make_shared<FieldSchema>("content", DataType::STRING, false,
fts_params),
std::make_shared<FieldSchema>("tag", DataType::INT32, false,
invert_params),
// Dummy vector field required for filter parsing path in
// execute
std::make_shared<FieldSchema>(
"vec", DataType::VECTOR_FP32, 4, false,
std::make_shared<FlatIndexParams>(MetricType::L2)),
});
}
static Segment::Ptr create_segment() {
auto segment_meta = std::make_shared<SegmentMeta>();
segment_meta->set_id(0);
auto id_map = IDMap::CreateAndOpen("fts_recall_test", seg_path_ + "/id_map",
true, false);
auto delete_store = std::make_shared<DeleteStore>("fts_recall_test");
Version v1;
v1.set_schema(*schema_);
std::string v_path = seg_path_ + "/manifest";
FileHelper::CreateDirectory(v_path);
auto vm = VersionManager::Create(v_path, v1);
if (!vm.has_value()) {
return nullptr;
}
BlockMeta mem_block;
mem_block.id_ = 0;
mem_block.type_ = BlockType::SCALAR;
mem_block.min_doc_id_ = 0;
mem_block.max_doc_id_ = 0;
mem_block.doc_count_ = 0;
segment_meta->set_writing_forward_block(mem_block);
SegmentOptions options;
options.read_only_ = false;
options.enable_mmap_ = true;
options.max_buffer_size_ = 256 * 1024;
auto result = Segment::CreateAndOpen(seg_path_, *schema_, 0, 0, id_map,
delete_store, vm.value(), options);
if (!result) {
return nullptr;
}
return result.value();
}
static void insert_docs(const Segment::Ptr &segment) {
// doc_id 0: "apple banana cherry" tag=1
// doc_id 1: "banana date elderberry" tag=2
// doc_id 2: "cherry fig grape" tag=1
// doc_id 3: "apple fig honeydew" tag=2
// doc_id 4: "date grape kiwi" tag=1
// doc_id 5: "apple apple apple" tag=2
// doc_id 6: "mango papaya starfruit" tag=1
// doc_id 7: "banana banana grape" tag=2
struct Entry {
std::string content;
int32_t tag;
};
std::vector<Entry> entries = {
{"apple banana cherry", 1}, {"banana date elderberry", 2},
{"cherry fig grape", 1}, {"apple fig honeydew", 2},
{"date grape kiwi", 1}, {"apple apple apple", 2},
{"mango papaya starfruit", 1}, {"banana banana grape", 2},
};
for (size_t i = 0; i < entries.size(); ++i) {
Doc doc;
doc.set_pk("pk_" + std::to_string(i));
doc.set_doc_id(i);
doc.set<std::string>("content", entries[i].content);
doc.set<int32_t>("tag", entries[i].tag);
auto status = segment->Insert(doc);
ASSERT_TRUE(status.ok())
<< "Insert doc " << i << " failed: " << status.c_str();
}
}
protected:
static inline std::string seg_path_ = "./fts_recall_test_collection";
static inline CollectionSchema::Ptr schema_;
static inline std::vector<Segment::Ptr> segments_;
static inline SQLEngine::Ptr engine_;
};
// ============================================================
// Basic FTS search tests
// ============================================================
// "apple" matches docs 0, 3, 5
TEST_F(FtsRecallTest, BasicSingleTerm) {
auto result = fts_search("apple");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 3u);
}
// BM25 ordering: doc 5 ("apple apple apple") should have highest score
TEST_F(FtsRecallTest, BM25ScoreOrdering) {
auto result = fts_search("apple");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
ASSERT_GE(result->size(), 2u);
// Results should be sorted by score descending
for (size_t i = 0; i + 1 < result->size(); ++i) {
EXPECT_GE((*result)[i]->score(), (*result)[i + 1]->score())
<< "Results not sorted descending at index " << i;
}
// Doc 5 has highest TF for "apple"
EXPECT_EQ((*result)[0]->pk(), "pk_5");
}
// "kiwi" only in doc 4
TEST_F(FtsRecallTest, SingleMatch) {
auto result = fts_search("kiwi");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
ASSERT_EQ(result->size(), 1u);
EXPECT_EQ((*result)[0]->pk(), "pk_4");
}
// Nonexistent term
TEST_F(FtsRecallTest, NoMatch) {
auto result = fts_search("zzznomatch");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 0u);
}
// Topk limit: "banana" in docs 0, 1, 7 (3 matches), topk=2
TEST_F(FtsRecallTest, TopkLimit) {
auto result = fts_search("banana", /*topk=*/2);
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_LE(result->size(), 2u);
}
// Multi-term implicit OR: "apple banana" matches union of {0,3,5} and {0,1,7}
TEST_F(FtsRecallTest, MultiTermImplicitOr) {
auto result = fts_search("apple banana");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
// Union: {0,1,3,5,7} = 5 docs
EXPECT_EQ(result->size(), 5u);
}
// "starfruit" only in doc 6
TEST_F(FtsRecallTest, RareTerm) {
auto result = fts_search("starfruit");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
ASSERT_EQ(result->size(), 1u);
EXPECT_EQ((*result)[0]->pk(), "pk_6");
}
// "grape" in docs 2, 4, 7
TEST_F(FtsRecallTest, CommonTerm) {
auto result = fts_search("grape");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 3u);
}
// ============================================================
// Explicit AND
// ============================================================
// "apple AND banana" -> intersection of {0,3,5} and {0,1,7} = {0}
TEST_F(FtsRecallTest, ExplicitAnd) {
auto result = fts_search("apple AND banana");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 1u);
EXPECT_EQ((*result)[0]->pk(), "pk_0");
}
// "cherry AND fig" -> {0,2} AND {2,3} = {2}
TEST_F(FtsRecallTest, ExplicitAnd2) {
auto result = fts_search("cherry AND fig");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 1u);
EXPECT_EQ((*result)[0]->pk(), "pk_2");
}
// ============================================================
// Binary NOT (AND-NOT)
// ============================================================
// "apple NOT banana" -> {0,3,5} minus {0,1,7} = {3,5}
TEST_F(FtsRecallTest, BinaryNot) {
auto result = fts_search("apple NOT banana");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 2u);
std::set<std::string> pks;
for (auto &doc : *result) {
pks.insert(doc->pk());
}
EXPECT_TRUE(pks.count("pk_3"));
EXPECT_TRUE(pks.count("pk_5"));
}
// "banana NOT grape" -> {0,1,7} minus {2,4,7} = {0,1}
TEST_F(FtsRecallTest, BinaryNot2) {
auto result = fts_search("banana NOT grape");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 2u);
std::set<std::string> pks;
for (auto &doc : *result) {
pks.insert(doc->pk());
}
EXPECT_TRUE(pks.count("pk_0"));
EXPECT_TRUE(pks.count("pk_1"));
}
// ============================================================
// Error cases
// ============================================================
// Leading NOT should fail parse
TEST_F(FtsRecallTest, LeadingNotIsRejected) {
auto result = fts_search("NOT apple");
EXPECT_FALSE(result.has_value());
}
// Both query_string_ and match_string_ empty
TEST_F(FtsRecallTest, BothEmptyReturnsError) {
SearchQuery vq;
vq.topk_ = 10;
vq.target_.field_name_ = "content";
vq.target_.clause_ = FtsClause{}; // both fields empty
auto result = engine_->execute(schema_, vq, segments_);
EXPECT_FALSE(result.has_value());
}
// Both query_string_ and match_string_ set
TEST_F(FtsRecallTest, BothSetReturnsError) {
SearchQuery vq;
vq.topk_ = 10;
vq.target_.field_name_ = "content";
FtsClause fts;
fts.query_string_ = "apple";
fts.match_string_ = "banana";
vq.target_.clause_ = fts;
auto result = engine_->execute(schema_, vq, segments_);
EXPECT_FALSE(result.has_value());
}
// ============================================================
// match_string tests
// ============================================================
// match_string "starfruit" -> doc 6
TEST_F(FtsRecallTest, MatchStringRareTerm) {
auto result = fts_match("starfruit");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
ASSERT_EQ(result->size(), 1u);
EXPECT_EQ((*result)[0]->pk(), "pk_6");
}
// match_string "grape" -> docs 2, 4, 7
TEST_F(FtsRecallTest, MatchStringCommonTerm) {
auto result = fts_match("grape");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 3u);
}
// match_string "apple banana" -> OR -> union {0,1,3,5,7}
TEST_F(FtsRecallTest, MatchStringMultipleTokens) {
auto result = fts_match("apple banana");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 5u);
}
// match_string analysing to zero tokens → empty result, not an error.
TEST_F(FtsRecallTest, MatchStringEmptyTokensReturnsNoResults) {
auto result = fts_match(" \t ");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_TRUE(result->empty());
}
// ============================================================
// default_operator tests
// ============================================================
// AND default for match_string: "apple banana" -> intersection = {0}
TEST_F(FtsRecallTest, DefaultOperatorAnd_MatchString) {
auto result = fts_match("apple banana", "AND");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 1u);
EXPECT_EQ((*result)[0]->pk(), "pk_0");
}
// OR default for match_string (backward compat)
TEST_F(FtsRecallTest, DefaultOperatorOr_MatchString) {
auto result = fts_match("apple banana", "OR");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 5u);
}
// AND default for query_string: "apple banana" -> AND
TEST_F(FtsRecallTest, DefaultOperatorAnd_QueryString) {
auto result = fts_query_with_op("apple banana", "AND");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 1u);
EXPECT_EQ((*result)[0]->pk(), "pk_0");
}
// Explicit OR in query not overridden by default_operator=AND
// "apple OR grape" with AND default -> OR still applies
TEST_F(FtsRecallTest, DefaultOperatorAnd_DoesNotOverrideExplicitOr) {
auto result = fts_query_with_op("apple OR grape", "AND");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
// apple: {0,3,5}, grape: {2,4,7} -> union = 6
EXPECT_EQ(result->size(), 6u);
}
// Empty default_operator keeps historical OR for match_string
TEST_F(FtsRecallTest, DefaultOperatorEmpty_BackwardCompatibleOr) {
auto result = fts_match("apple banana"); // no default_op arg
ASSERT_TRUE(result.has_value()) << result.error().c_str();
// OR semantics: union of apple{0,3,5} and banana{0,1,7} = 5
EXPECT_EQ(result->size(), 5u);
}
// Lowercase "and" must be accepted
TEST_F(FtsRecallTest, DefaultOperatorAndLowercase_Accepted) {
auto result = fts_match("apple banana", "and");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 1u);
}
// Mixed-case "And" / "oR" are accepted via case-insensitive normalisation.
TEST_F(FtsRecallTest, DefaultOperatorMixedCase_Accepted) {
{
// "And" -> AND semantics: intersection of apple{0,3,5} and banana{0,1,7}
auto result = fts_match("apple banana", "And");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 1u);
}
{
// "oR" -> OR semantics: union = 5 docs
auto result = fts_match("apple banana", "oR");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_EQ(result->size(), 5u);
}
}
// Invalid default_operator value should be rejected (was previously silently
// downgraded to OR).
TEST_F(FtsRecallTest, DefaultOperatorInvalid_Rejected) {
auto result = fts_match("apple banana", "xor");
EXPECT_FALSE(result.has_value());
}
// ============================================================
// Error cases (additional)
// ============================================================
// Empty field_name should fail
TEST_F(FtsRecallTest, EmptyFieldNameReturnsError) {
SearchQuery vq;
vq.topk_ = 10;
vq.target_.field_name_ = "";
FtsClause fts;
fts.query_string_ = "apple";
vq.target_.clause_ = fts;
auto result = engine_->execute(schema_, vq, segments_);
EXPECT_FALSE(result.has_value());
}
// Empty query_string (with field_name set) should fail
TEST_F(FtsRecallTest, EmptyQueryStringReturnsError) {
SearchQuery vq;
vq.topk_ = 10;
vq.target_.field_name_ = "content";
// Both query_string_ and match_string_ empty -> error
vq.target_.clause_ = FtsClause{};
auto result = engine_->execute(schema_, vq, segments_);
EXPECT_FALSE(result.has_value());
}
// ============================================================
// FTS search with WHERE filter
// ============================================================
// "apple" (docs 0,3,5) + tag = 1 (docs 0,2,4,6) -> intersection = {0}
TEST_F(FtsRecallTest, FtsSearchWithFilter_ScoreTag) {
auto result = fts_search_with_filter("apple", "tag = 1");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
// Filter should reduce results to doc 0 only
EXPECT_LE(result->size(), 3u);
// Verify that at least doc 0 (which satisfies both FTS and filter) is present
bool found_pk0 = false;
for (auto &doc : *result) {
if (doc->pk() == "pk_0") {
found_pk0 = true;
}
}
EXPECT_TRUE(found_pk0);
}
// "banana" (docs 0,1,7) + tag = 2 (docs 1,3,5,7) + topk=1
TEST_F(FtsRecallTest, FtsSearchWithFilter_TopkRespected) {
auto result = fts_search_with_filter("banana", "tag = 2", /*topk=*/1);
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_LE(result->size(), 1u);
}
// "apple" matches docs 0,3,5, but no doc has tag=999.
TEST_F(FtsRecallTest, FtsSearchWithFilter_ZeroMatchesReturnsEmpty) {
auto result = fts_search_with_filter("apple", "tag = 999");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_TRUE(result->empty());
}
// An FTS field can only be used as a query target, not as a filter condition.
// Putting the FTS field ("content") in the WHERE filter must be rejected.
TEST_F(FtsRecallTest, FtsFieldNotAllowedInFilter) {
auto result = fts_search_with_filter("apple", "content = 'apple'");
ASSERT_FALSE(result.has_value());
}
// ============================================================
// Repeated-term linearity: the AST rewriter collapses a repeated term into a
// single TermNode whose boost equals the occurrence count. With linear boost
// the per-document score must be exactly N× the single-term score, matching
// the pre-rewrite "N independent scorers summed" semantics.
// ============================================================
TEST_F(FtsRecallTest, MatchStringRepeatedTermLinearBoost) {
auto baseline = fts_match("apple");
auto repeated = fts_match("apple apple");
ASSERT_TRUE(baseline.has_value()) << baseline.error().c_str();
ASSERT_TRUE(repeated.has_value()) << repeated.error().c_str();
ASSERT_EQ(baseline->size(), repeated->size());
// Same doc set, same ordering — only the absolute scores differ.
for (size_t i = 0; i < baseline->size(); ++i) {
EXPECT_EQ((*baseline)[i]->pk(), (*repeated)[i]->pk()) << "rank " << i;
EXPECT_FLOAT_EQ((*baseline)[i]->score() * 2.0f, (*repeated)[i]->score())
<< "rank " << i << " pk=" << (*repeated)[i]->pk();
}
}
// Unary `-` prefix inside an OR was previously executed via build_or_iterator
// wrapping the disjunction in a must_not Conjunction. After the rewriter
// canonicalizes OR-with-must_not into AND(positive..., -negative...), the
// must_not iterator path lives only in build_and_iterator. End-to-end the
// match set must be unchanged: apple{0,3,5} banana{0,1,7} = {3, 5}.
TEST_F(FtsRecallTest, QueryStringUnaryMinusExcludesMatchingDocs) {
auto result = fts_search("apple -banana");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
std::set<std::string> pks;
for (const auto &d : *result) {
pks.insert(d->pk());
}
EXPECT_EQ(pks, std::set<std::string>({"pk_3", "pk_5"}));
}
// `apple -apple` is a self-contradiction; the rewriter detects the must vs
// must_not conflict after canonicalization and rewrites the whole subtree
// to EmptyNode, so the query short-circuits to zero docs.
TEST_F(FtsRecallTest, QueryStringSelfContradictionReturnsNoResults) {
auto result = fts_search("apple -apple");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_TRUE(result->empty());
}
TEST_F(FtsRecallTest, MatchStringRepeatedTermPreservesUnion) {
// "apple apple banana" — apple repeated, banana once. Doc set must equal
// "apple banana" (union), and apple-only docs should score 2× their
// single-term score plus zero for banana.
auto plain_union = fts_match("apple banana");
auto repeated_union = fts_match("apple apple banana");
ASSERT_TRUE(plain_union.has_value()) << plain_union.error().c_str();
ASSERT_TRUE(repeated_union.has_value()) << repeated_union.error().c_str();
EXPECT_EQ(plain_union->size(), repeated_union->size());
std::set<std::string> plain_pks;
std::set<std::string> repeated_pks;
for (const auto &d : *plain_union) {
plain_pks.insert(d->pk());
}
for (const auto &d : *repeated_union) {
repeated_pks.insert(d->pk());
}
EXPECT_EQ(plain_pks, repeated_pks);
}
// ============================================================
// FTS delete / upsert end-to-end tests (per-test fixture)
// ============================================================
class FtsRecallDeleteTest : public ::testing::Test {
protected:
void SetUp() override {
seg_path_ = "./fts_recall_delete_test_" +
std::to_string(reinterpret_cast<uintptr_t>(this));
FileHelper::RemoveDirectory(seg_path_);
FileHelper::CreateDirectory(seg_path_);
auto fts_params = std::make_shared<FtsIndexParams>(
"whitespace", std::vector<std::string>{"lowercase"}, "");
auto invert_params = std::make_shared<InvertIndexParams>(true);
schema_ = std::make_shared<CollectionSchema>(
"fts_delete_test",
std::vector<FieldSchema::Ptr>{
std::make_shared<FieldSchema>("content", DataType::STRING, false,
fts_params),
std::make_shared<FieldSchema>("tag", DataType::INT32, false,
invert_params),
std::make_shared<FieldSchema>(
"vec", DataType::VECTOR_FP32, 4, false,
std::make_shared<FlatIndexParams>(MetricType::L2)),
});
auto segment_meta = std::make_shared<SegmentMeta>();
segment_meta->set_id(0);
auto id_map = IDMap::CreateAndOpen("fts_delete_test", seg_path_ + "/id_map",
true, false);
auto delete_store = std::make_shared<DeleteStore>("fts_delete_test");
Version v1;
v1.set_schema(*schema_);
std::string v_path = seg_path_ + "/manifest";
FileHelper::CreateDirectory(v_path);
auto vm = VersionManager::Create(v_path, v1);
ASSERT_TRUE(vm.has_value());
BlockMeta mem_block;
mem_block.id_ = 0;
mem_block.type_ = BlockType::SCALAR;
mem_block.min_doc_id_ = 0;
mem_block.max_doc_id_ = 0;
mem_block.doc_count_ = 0;
segment_meta->set_writing_forward_block(mem_block);
SegmentOptions options;
options.read_only_ = false;
options.enable_mmap_ = true;
options.max_buffer_size_ = 256 * 1024;
auto result = Segment::CreateAndOpen(seg_path_, *schema_, 0, 0, id_map,
delete_store, vm.value(), options);
ASSERT_TRUE(result.has_value());
segment_ = result.value();
segments_.push_back(segment_);
engine_ = SQLEngine::create(std::make_shared<Profiler>());
insert_docs();
}
void TearDown() override {
segments_.clear();
segment_.reset();
engine_.reset();
schema_.reset();
FileHelper::RemoveDirectory(seg_path_);
}
void insert_docs() {
// doc_id 0: "apple banana cherry" tag=1
// doc_id 1: "banana date elderberry" tag=2
// doc_id 2: "cherry fig grape" tag=1
// doc_id 3: "apple fig honeydew" tag=2
// doc_id 4: "date grape kiwi" tag=1
struct Entry {
std::string content;
int32_t tag;
};
std::vector<Entry> entries = {
{"apple banana cherry", 1}, {"banana date elderberry", 2},
{"cherry fig grape", 1}, {"apple fig honeydew", 2},
{"date grape kiwi", 1},
};
for (size_t i = 0; i < entries.size(); ++i) {
Doc doc;
doc.set_pk("pk_" + std::to_string(i));
doc.set_doc_id(i);
doc.set<std::string>("content", entries[i].content);
doc.set<int32_t>("tag", entries[i].tag);
auto status = segment_->Insert(doc);
ASSERT_TRUE(status.ok())
<< "Insert doc " << i << " failed: " << status.c_str();
}
}
Result<DocPtrList> fts_search(const std::string &query_string,
int topk = 10) {
SearchQuery vq;
vq.topk_ = topk;
vq.target_.field_name_ = "content";
FtsClause fts;
fts.query_string_ = query_string;
vq.target_.clause_ = fts;
return engine_->execute(schema_, vq, segments_);
}
std::set<std::string> collect_pks(const DocPtrList &docs) {
std::set<std::string> pks;
for (const auto &d : docs) {
pks.insert(d->pk());
}
return pks;
}
std::string seg_path_;
CollectionSchema::Ptr schema_;
Segment::Ptr segment_;
std::vector<Segment::Ptr> segments_;
SQLEngine::Ptr engine_;
};
// Delete doc 0 ("apple banana cherry"), then search "apple":
// before: {0, 3}, after: {3} only.
TEST_F(FtsRecallDeleteTest, DeletedDocExcludedFromSearch) {
auto before = fts_search("apple");
ASSERT_TRUE(before.has_value()) << before.error().c_str();
EXPECT_TRUE(collect_pks(*before).count("pk_0"));
auto s = segment_->Delete("pk_0");
ASSERT_TRUE(s.ok()) << s.c_str();
auto after = fts_search("apple");
ASSERT_TRUE(after.has_value()) << after.error().c_str();
auto pks = collect_pks(*after);
EXPECT_FALSE(pks.count("pk_0"));
EXPECT_TRUE(pks.count("pk_3"));
}
// Delete all docs matching "banana" (0, 1), verify "banana" returns empty.
TEST_F(FtsRecallDeleteTest, DeleteAllMatchingDocsReturnsEmpty) {
auto s1 = segment_->Delete("pk_0");
ASSERT_TRUE(s1.ok()) << s1.c_str();
auto s2 = segment_->Delete("pk_1");
ASSERT_TRUE(s2.ok()) << s2.c_str();
auto result = fts_search("banana");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_TRUE(result->empty());
}
// Upsert doc 0 with new content, verify old content no longer matches
// and new content is searchable.
TEST_F(FtsRecallDeleteTest, UpsertUpdatesSearchableContent) {
// Before: "apple" matches {0, 3}
auto before = fts_search("apple");
ASSERT_TRUE(before.has_value()) << before.error().c_str();
EXPECT_EQ(before->size(), 2u);
// Upsert pk_0 with completely different content
Doc updated;
updated.set_pk("pk_0");
updated.set<std::string>("content", "mango pineapple watermelon");
updated.set<int32_t>("tag", 1);
auto s = segment_->Upsert(updated);
ASSERT_TRUE(s.ok()) << s.c_str();
// "apple" should now only match doc 3
auto after_apple = fts_search("apple");
ASSERT_TRUE(after_apple.has_value()) << after_apple.error().c_str();
ASSERT_EQ(after_apple->size(), 1u);
EXPECT_EQ((*after_apple)[0]->pk(), "pk_3");
// "pineapple" should match the upserted doc
auto after_new = fts_search("pineapple");
ASSERT_TRUE(after_new.has_value()) << after_new.error().c_str();
ASSERT_EQ(after_new->size(), 1u);
EXPECT_EQ((*after_new)[0]->pk(), "pk_0");
}
// Delete a doc, then search with AND: "cherry AND fig" was {2},
// delete doc 2 → empty.
TEST_F(FtsRecallDeleteTest, DeleteAffectsConjunctionQuery) {
auto before = fts_search("cherry AND fig");
ASSERT_TRUE(before.has_value()) << before.error().c_str();
ASSERT_EQ(before->size(), 1u);
EXPECT_EQ((*before)[0]->pk(), "pk_2");
auto s = segment_->Delete("pk_2");
ASSERT_TRUE(s.ok()) << s.c_str();
auto after = fts_search("cherry AND fig");
ASSERT_TRUE(after.has_value()) << after.error().c_str();
EXPECT_TRUE(after->empty());
}
// Delete a doc, flush, then verify deleted doc stays excluded.
TEST_F(FtsRecallDeleteTest, DeletePersistsAcrossFlush) {
auto s = segment_->Delete("pk_4");
ASSERT_TRUE(s.ok()) << s.c_str();
auto flush_s = segment_->flush();
ASSERT_TRUE(flush_s.ok()) << flush_s.c_str();
auto result = fts_search("kiwi");
ASSERT_TRUE(result.has_value()) << result.error().c_str();
EXPECT_TRUE(result->empty());
}
} // namespace zvec::sqlengine
+789
View File
@@ -0,0 +1,789 @@
// 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 <cstdint>
#include <memory>
#include <gtest/gtest.h>
#include "db/sqlengine/sqlengine.h"
#include "zvec/db/schema.h"
#include "recall_base.h"
namespace zvec::sqlengine {
class InvertRecallTest : public RecallTest {};
TEST_F(InvertRecallTest, Eq) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_age = 1";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 100);
for (int j = 0, i = 1; j < (int)docs.size(); j++, i += 100) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(InvertRecallTest, Gt) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_id > 1000";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int j = 0; j < query.topk_; j++) {
auto &doc = docs[j];
auto i = j + 1001;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(InvertRecallTest, Ge) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_id >= 1000";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int j = 0; j < query.topk_; j++) {
auto &doc = docs[j];
auto i = j + 1000;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(InvertRecallTest, Lt) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_id < 100";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
ASSERT_EQ(docs.size(), 100);
for (int j = 0, i = 0; j < (int)docs.size(); j++, i += 1) {
auto &doc = docs[j];
EXPECT_EQ(i, doc->get<uint64_t>("id"));
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(InvertRecallTest, Le) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_id <= 100";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
ASSERT_EQ(docs.size(), 101);
for (int j = 0, i = 0; j < (int)docs.size(); j++, i += 1) {
auto &doc = docs[j];
EXPECT_EQ(i, doc->get<uint64_t>("id"));
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(InvertRecallTest, And) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_id <= 100 and invert_id > 50";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
ASSERT_EQ(docs.size(), 50);
for (int j = 0, i = 51; j < (int)docs.size(); j++, i += 1) {
auto &doc = docs[j];
EXPECT_EQ(i, doc->get<uint64_t>("id"));
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(InvertRecallTest, Or) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_id < 100 or invert_id > 200";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
ASSERT_EQ(docs.size(), 200);
for (int j = 0; j < (int)docs.size(); j++) {
int i = j < 100 ? j : j + 101;
auto &doc = docs[j];
EXPECT_EQ(i, doc->get<uint64_t>("id"));
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(InvertRecallTest, StrEq) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_name = 'user_1'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 100);
for (int j = 0, i = 1; j < (int)docs.size(); j++, i += 100) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(InvertRecallTest, StrGe) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_name >= 'user_1'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 0; j < (int)docs.size(); j++, i += 1) {
if (i % 100 == 0) {
i += 1;
}
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(InvertRecallTest, StrIn) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_name IN ('user_1', 'user_2')";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 1; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
if (i % 100 == 1) {
i += 1;
} else if (i % 100 == 2) {
i += 99;
}
}
}
TEST_F(InvertRecallTest, StrNotIn) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_name NOT IN ('user_1', 'user_2')";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 0; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
if (i % 100 == 0) {
i += 3;
} else {
i += 1;
}
}
}
TEST_F(InvertRecallTest, StrLike) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_name like 'user\\_9%'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 9; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
if (i % 100 == 9) {
i += 81;
} else if (i % 100 == 99) {
i += 10;
} else {
i += 1;
}
}
}
TEST_F(InvertRecallTest, ContainAll) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_category_set contain_all (";
for (int i = 1; i <= 32; i++) {
query.filter_ += std::to_string(i);
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
while (i % 100 < 32) {
i += 1;
}
}
}
TEST_F(InvertRecallTest, NotContainAll) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_category_set not contain_all (";
for (int i = 1; i <= 32; i++) {
query.filter_ += std::to_string(i);
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 1; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
// i % 100 == 0 has null category
while (i % 100 >= 32 || i % 100 == 0) {
i += 1;
}
}
}
TEST_F(InvertRecallTest, ContainAny) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_category_set contain_any (98,99,100)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 98; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
while (i % 100 < 98) {
i += 1;
}
}
}
TEST_F(InvertRecallTest, NotContainAny) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_category_set not contain_any (98,99,100)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 1; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
// i % 100 == 0 has null category
while (i % 100 >= 98 || i % 100 == 0) {
i += 1;
}
}
}
TEST_F(InvertRecallTest, BoolContainAll) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_bool_array contain_all (true, false)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 0; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 3;
}
}
TEST_F(InvertRecallTest, BoolContainAny) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_bool_array contain_any (true)";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 0; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
if (i % 3 == 2) {
i += 1;
}
}
}
TEST_F(InvertRecallTest, ContainAllEmptySet) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_category_set contain_all ()";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 1; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
while (i % 100 == 0) {
i += 1;
}
}
}
TEST_F(InvertRecallTest, NotContainAllEmptySet) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_category_set not contain_all ()";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
ASSERT_EQ(docs.size(), 0);
}
TEST_F(InvertRecallTest, ContainAnyEmptySet) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_category_set contain_any ()";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
ASSERT_EQ(docs.size(), 0);
}
TEST_F(InvertRecallTest, NotContainAnyEmptySet) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_category_set not contain_any ()";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 1; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
while (i % 100 == 0) {
i += 1;
}
}
}
TEST_F(InvertRecallTest, IsNull) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_optional_age is null";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 100);
for (int j = 0, i = 0; j < (int)docs.size(); j++, i += 100) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(InvertRecallTest, IsNotNull) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_optional_age is not null";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 0; j < (int)docs.size(); j++, i += 1) {
if (i % 100 == 0) {
i += 1;
}
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(InvertRecallTest, BoolEqTrue) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_bool = TRuE";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 100);
for (int j = 0, i = 0; j < (int)docs.size(); j++, i += 100) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
}
}
TEST_F(InvertRecallTest, BoolEqFalse) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "invert_bool = false";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 1; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
if (i % 100 == 0) {
i += 1;
}
}
}
TEST_F(InvertRecallTest, ArrayLengthGe) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "array_length(invert_category_set) >= 32";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 200);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 1;
while (i % 100 < 32) {
i += 1;
}
}
}
TEST_F(InvertRecallTest, ArrayLengthEq) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
query.filter_ = "array_length(invert_category_set) = 32";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), 100);
for (int j = 0, i = 32; j < (int)docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
EXPECT_EQ(i, doc->get<uint64_t>("id"));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
i += 100;
}
}
TEST_F(InvertRecallTest, MultiSegment) {
SearchQuery query;
query.output_fields_ = std::vector<std::string>();
query.topk_ = 200;
query.include_vector_ = true;
query.filter_ = "invert_id <= 5000";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
std::vector<Segment::Ptr> segments = segments_;
segments.push_back(segments_[0]);
auto ret = engine->execute(collection_schema_, query, segments);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int i = 0; i < query.topk_; i++) {
auto &doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto dense = doc->get<std::vector<float>>("dense");
ASSERT_TRUE(dense.has_value());
EXPECT_EQ(dense.value().size(), 4);
for (auto v : dense.value()) {
EXPECT_FLOAT_EQ(v, (float)i);
}
auto sparse =
doc->get<std::pair<std::vector<uint32_t>, std::vector<float>>>(
"sparse");
if (i % 100 == 0) {
// set with empty vector
ASSERT_FALSE(sparse.has_value());
continue;
}
ASSERT_TRUE(sparse.has_value());
const auto &[indices, values] = sparse.value();
EXPECT_EQ(indices.size(), i % 100);
EXPECT_EQ(values.size(), i % 100);
for (int j = 0; j < i % 100; j++) {
EXPECT_EQ(indices[j], j);
EXPECT_FLOAT_EQ(values[j], (float)i);
}
}
}
} // namespace zvec::sqlengine
+372
View File
@@ -0,0 +1,372 @@
// 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 <cstdint>
#include <cstdlib>
#include <iostream>
#include <memory>
#include <arrow/api.h>
#include <arrow/io/api.h>
#include <arrow/ipc/api.h>
#include <gtest/gtest.h>
#include "db/common/file_helper.h"
#include "db/index/common/version_manager.h"
#include "db/index/segment/segment.h"
#include "db/sqlengine/sqlengine.h"
#include "zvec/db/index_params.h"
#include "zvec/db/schema.h"
#include "zvec/db/type.h"
#include "test_helper.h"
namespace zvec::sqlengine {
static Doc create_doc(const uint64_t doc_id) {
Doc new_doc;
new_doc.set_pk("pk_" + std::to_string(doc_id));
new_doc.set_doc_id(doc_id);
auto name = std::string("user-");
if (doc_id >= 5000 && doc_id < 8000) {
name += "%";
} else if (doc_id >= 8000) {
name += '_';
}
name += std::to_string(doc_id % 100);
new_doc.set<std::string>("name", name);
new_doc.set<std::string>("invert_name", name);
new_doc.set<std::string>("extended_invert_name", name);
return new_doc;
}
class LikeTest : public testing::Test {
protected:
static void SetUpTestSuite() {
FileHelper::RemoveDirectory(seg_path_);
FileHelper::CreateDirectory(seg_path_);
auto invert_params = std::make_shared<InvertIndexParams>(true);
collection_schema_ = std::make_shared<CollectionSchema>(
"test_collection",
std::vector<FieldSchema::Ptr>{
std::make_shared<FieldSchema>("name", DataType::STRING, false,
nullptr),
std::make_shared<FieldSchema>(
"invert_name", DataType::STRING, false,
std::make_shared<InvertIndexParams>(false, false)),
std::make_shared<FieldSchema>(
"extended_invert_name", DataType::STRING, false,
std::make_shared<InvertIndexParams>(false, true)),
});
auto segment = create_segment(seg_path_, *collection_schema_);
if (segment == nullptr) {
LOG_ERROR("create segment failed");
EXPECT_TRUE(segment != nullptr);
std::exit(EXIT_FAILURE);
}
auto status = InsertDoc(segment, 0, 10000, &create_doc);
if (!status.ok()) {
LOG_ERROR("insert doc failed: %s", status.c_str());
EXPECT_TRUE(status.ok());
std::exit(EXIT_FAILURE);
}
segments_.push_back(segment);
}
static void TearDownTestSuite() {
segments_.clear();
FileHelper::RemoveDirectory(seg_path_);
}
protected:
static inline std::string seg_path_ = "./test_collection";
static inline CollectionSchema::Ptr collection_schema_;
static inline std::vector<Segment::Ptr> segments_;
};
TEST_F(LikeTest, ForwardLikeAll) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = "name like '%'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0; i < docs.size(); i++) {
auto doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
}
}
TEST_F(LikeTest, InvertLikeAll) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = "invert_name like '%'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0; i < docs.size(); i++) {
auto doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
}
}
TEST_F(LikeTest, ForwardPrefixLike) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = "name like 'user-22%'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0; i < docs.size(); i++) {
auto doc = docs[i];
int doc_id = i * 100 + 22;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, InvertPrefixLike) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = "invert_name like 'user-22%'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0; i < docs.size(); i++) {
auto doc = docs[i];
int doc_id = i * 100 + 22;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, ForwardSuffixLike) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = "name like '%ser-22'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0; i < docs.size(); i++) {
auto doc = docs[i];
int doc_id = i * 100 + 22;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, NotExtendedInvertSuffixLikeRunAsForward) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = "invert_name like '%ser-22'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0; i < docs.size(); i++) {
auto doc = docs[i];
int doc_id = i * 100 + 22;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, ExtendedInvertSuffixLike) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = "extended_invert_name like '%ser-22'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0; i < docs.size(); i++) {
auto doc = docs[i];
int doc_id = i * 100 + 22;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, ForwardMiddleLike) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = "name like 'user%2'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0, doc_id = 0; i < docs.size(); i++, doc_id++) {
auto doc = docs[i];
while (doc_id % 100 % 10 != 2) {
doc_id++;
}
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, ExtendedInvertMiddleLike) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = "extended_invert_name like 'user%2'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0, doc_id = 0; i < docs.size(); i++, doc_id++) {
auto doc = docs[i];
while (doc_id % 100 % 10 != 2) {
doc_id++;
}
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, UnderScore) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = "name like 'user-_2'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0, doc_id = 0; i < docs.size(); i++, doc_id++) {
auto doc = docs[i];
while (doc_id % 100 % 10 != 2 || doc_id % 100 < 10) {
doc_id++;
}
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, InvertUnderScoreRunAsForward) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = "invert_name like 'user-_2'";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0, doc_id = 0; i < docs.size(); i++, doc_id++) {
auto doc = docs[i];
while (doc_id % 100 % 10 != 2 || doc_id % 100 < 10) {
doc_id++;
}
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, ForwardEscapePercent) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = R"(name like 'user-\%%')";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0, doc_id = 5000; i < docs.size(); i++, doc_id++) {
auto doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, InvertEscapePercent) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = R"(invert_name like 'user-\%%')";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0, doc_id = 5000; i < docs.size(); i++, doc_id++) {
auto doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, ForwardEscapeUnderscore) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = R"(name like 'user-\_%')";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0, doc_id = 8000; i < docs.size(); i++, doc_id++) {
auto doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, InvertEscapeUnderscore) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = R"(invert_name like 'user-\_%')";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
ASSERT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0, doc_id = 8000; i < docs.size(); i++, doc_id++) {
auto doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
TEST_F(LikeTest, NoPercentRunAsEqual) {
SearchQuery query;
query.output_fields_ = {"name"};
query.topk_ = 200;
query.filter_ = R"(invert_name like 'user-22')";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
EXPECT_TRUE(ret.has_value()) << ret.error();
auto docs = std::move(ret.value());
for (size_t i = 0; i < docs.size(); i++) {
auto doc = docs[i];
int doc_id = i * 100 + 22;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
}
}
} // namespace zvec::sqlengine
+545
View File
@@ -0,0 +1,545 @@
// 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.
#pragma once
#include <memory>
#include <string>
#include <vector>
#include <arrow/api.h>
#include <arrow/compute/api.h>
#include <gtest/gtest.h>
#include <zvec/ailego/logger/logger.h>
#include "db/index/column/common/index_results.h"
#include "db/index/column/vector_column/vector_column_indexer.h"
#include "db/index/segment/segment.h"
#include "index/column/inverted_column/inverted_column_indexer.h"
#include "index/column/vector_column/vector_column_params.h"
#include "index/common/index_filter.h"
namespace zvec {
class MockIndexResult : public InvertedSearchResult {
public:
MockIndexResult(const std::vector<idx_t> &doc_ids,
const std::vector<float> &scores)
: doc_ids_(doc_ids), scores_(scores) {}
MockIndexResult(const std::vector<idx_t> &doc_ids,
const std::vector<float> &scores,
const std::vector<std::string> &groups)
: doc_ids_(doc_ids), scores_(scores), group_ids_(groups) {}
size_t count() const override {
return doc_ids_.size();
}
IteratorUPtr create_iterator() override {
return std::make_unique<MockIterator>(*this);
}
private:
struct MockIterator : public IndexResults::Iterator {
MockIterator(MockIndexResult &parent) : parent_(parent) {}
idx_t doc_id() const override {
return parent_.doc_ids_[current_index_];
}
float score() const override {
return parent_.scores_[current_index_];
}
void next() override {
++current_index_;
}
bool valid() const override {
return current_index_ < parent_.count();
}
const std::string &group_id() const override {
return parent_.group_ids_[current_index_];
}
MockIndexResult &parent_;
size_t current_index_{0};
};
std::vector<idx_t> doc_ids_;
std::vector<float> scores_;
std::vector<std::string> group_ids_;
};
class MockVectorIndexer : public CombinedVectorColumnIndexer {
public:
//! Search results with query
Result<IndexResults::Ptr> Search(
const vector_column_params::VectorData &vector_data,
const vector_column_params::QueryParams &query_params) override {
// return tl::make_unexpected(Status::InternalError("err"));
return std::make_shared<MockIndexResult>(
std::vector<idx_t>{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
std::vector<float>{0.1F, 0.2F, 0.3F, 0.4F, 0.5F, 0.6F, 0.7F, 0.8F, 0.9F,
1.0F},
std::vector<std::string>{"group_0", "group_1", "group_2", "group_0",
"group_1", "group_2", "group_0", "group_1",
"group_2", "group_0"});
}
Result<vector_column_params::VectorDataBuffer> Fetch(
uint32_t doc_id) const override {
// float f = doc_id;
// std::vector<float> v(4, f);
// std::string v_str = std::string(reinterpret_cast<char *>(v.data()),
// v.size() * sizeof(float));
// return vector_column_params::VectorDataBuffer{
// vector_column_params::DenseVectorBuffer{v_str}};
// sparse
uint32_t count = doc_id % 5;
std::vector<uint32_t> indices(count);
std::vector<float> values(count);
for (uint32_t i = 0; i < count; i++) {
indices[i] = i;
values[i] = i / 100.0;
}
return vector_column_params::VectorDataBuffer{
vector_column_params::SparseVectorBuffer{
std::string(reinterpret_cast<char *>(indices.data()),
indices.size() * sizeof(uint32_t)),
std::string(reinterpret_cast<char *>(values.data()),
values.size() * sizeof(float))}};
}
};
class MockInvertIndexer : public InvertedColumnIndexer {
public:
MockInvertIndexer() : InvertedColumnIndexer(ctx) {}
InvertedSearchResult::Ptr search(const std::string &value,
CompareOp op) const override {
return std::make_shared<MockIndexResult>(
std::vector<idx_t>{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
std::vector<float>{0.1F, 0.2F, 0.3F, 0.4F, 0.5F, 0.6F, 0.7F, 0.8F, 0.9F,
1.0F});
}
InvertedSearchResult::Ptr search_null() const override {
return std::make_shared<MockIndexResult>(
std::vector<idx_t>{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
std::vector<float>{0.1F, 0.2F, 0.3F, 0.4F, 0.5F, 0.6F, 0.7F, 0.8F, 0.9F,
1.0F});
}
InvertedSearchResult::Ptr search_non_null() const override {
return std::make_shared<MockIndexResult>(
std::vector<idx_t>{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
std::vector<float>{0.1F, 0.2F, 0.3F, 0.4F, 0.5F, 0.6F, 0.7F, 0.8F, 0.9F,
1.0F});
}
private:
RocksdbContext ctx;
};
// std::make_shared<FieldSchema>("id", DataType::INT32, false, 0, false,
// nullptr),
// std::make_shared<FieldSchema>("name", DataType::STRING, false, 0,
// false, nullptr),
// std::make_shared<FieldSchema>("age", DataType::INT64, false, 0,
// false, nullptr),
// std::make_shared<FieldSchema>("score", DataType::DOUBLE, false, 0,
// false, nullptr),
inline arrow::Result<std::shared_ptr<arrow::Table>> CreateTable(
int count = 10000000) {
auto schema = arrow::schema({
arrow::field("id", arrow::int32()),
arrow::field("name", arrow::utf8()),
arrow::field("age", arrow::int64()),
arrow::field("score", arrow::float64()),
arrow::field("_zvec_uid_", arrow::utf8()),
arrow::field("_zvec_row_id_", arrow::uint64()),
arrow::field("_zvec_g_doc_id_", arrow::uint64()),
arrow::field("tag_list", arrow::list(arrow::int32())),
});
std::shared_ptr<arrow::Array> array_id;
std::shared_ptr<arrow::Array> array_name;
std::shared_ptr<arrow::Array> array_age;
std::shared_ptr<arrow::Array> array_score;
std::shared_ptr<arrow::Array> array_uid;
arrow::NumericBuilder<arrow::Int64Type> builder;
auto has_value = [](int i) { return i % 13 != 0; };
ARROW_RETURN_NOT_OK(builder.Reserve(count));
for (int i = 0; i < count; i++) {
if (has_value(i)) {
ARROW_RETURN_NOT_OK((builder.Append(i)));
} else {
ARROW_RETURN_NOT_OK((builder.AppendNull()));
}
}
ARROW_RETURN_NOT_OK(builder.Finish(&array_age));
builder.Reset();
arrow::NumericBuilder<arrow::Int32Type> builder_id;
ARROW_RETURN_NOT_OK(builder_id.Reserve(count));
for (int i = 0; i < count; i++) {
if (has_value(i)) {
ARROW_RETURN_NOT_OK((builder_id.Append(i)));
} else {
ARROW_RETURN_NOT_OK((builder_id.AppendNull()));
}
}
ARROW_RETURN_NOT_OK(builder_id.Finish(&array_id));
arrow::NumericBuilder<arrow::DoubleType> builder_score;
ARROW_RETURN_NOT_OK(builder_score.Reserve(count));
for (int i = 0; i < count; i++) {
if (has_value(i)) {
ARROW_RETURN_NOT_OK((builder_score.Append(i / 100.0)));
} else {
ARROW_RETURN_NOT_OK((builder_score.AppendNull()));
}
}
ARROW_RETURN_NOT_OK(builder_score.Finish(&array_score));
arrow::StringBuilder builder_d;
ARROW_RETURN_NOT_OK(builder_d.Reserve(count));
for (int i = 0; i < count; i++) {
if (has_value(i)) {
ARROW_RETURN_NOT_OK((builder_d.Append("name_" + std::to_string(i))));
} else {
ARROW_RETURN_NOT_OK((builder_d.AppendNull()));
}
}
ARROW_RETURN_NOT_OK(builder_d.Finish(&array_name));
arrow::StringBuilder builder_uid;
ARROW_RETURN_NOT_OK(builder_uid.Reserve(count));
for (int i = 0; i < count; i++) {
ARROW_RETURN_NOT_OK((builder_uid.Append("uid_" + std::to_string(i))));
}
ARROW_RETURN_NOT_OK(builder_uid.Finish(&array_uid));
arrow::NumericBuilder<arrow::UInt64Type> builder_row_id;
ARROW_RETURN_NOT_OK(builder_row_id.Reserve(count));
for (int i = 0; i < count; i++) {
ARROW_RETURN_NOT_OK((builder_row_id.Append(i)));
}
std::shared_ptr<arrow::Array> array_row_id;
ARROW_RETURN_NOT_OK(builder_row_id.Finish(&array_row_id));
arrow::NumericBuilder<arrow::UInt64Type> builder_doc_id;
ARROW_RETURN_NOT_OK(builder_doc_id.Reserve(count));
for (int i = 0; i < count; i++) {
ARROW_RETURN_NOT_OK((builder_doc_id.Append(i)));
}
std::shared_ptr<arrow::Array> array_doc_id;
ARROW_RETURN_NOT_OK(builder_doc_id.Finish(&array_doc_id));
arrow::ListBuilder list_builder(arrow::default_memory_pool(),
std::make_shared<arrow::Int32Builder>());
auto *tag_value_builder =
static_cast<arrow::Int32Builder *>(list_builder.value_builder());
for (int i = 0; i < count; ++i) {
// 开始一个新的 list
ARROW_RETURN_NOT_OK(list_builder.Append());
int idx = i % 5; // 对应模式
for (int j = 0; j < idx + 1; ++j) {
ARROW_RETURN_NOT_OK(tag_value_builder->Append(j + 1));
}
}
std::shared_ptr<arrow::Array> tag_list_array;
auto status = list_builder.Finish(&tag_list_array);
;
return arrow::Table::Make(
schema, {array_id, array_name, array_age, array_score, array_uid,
array_row_id, array_doc_id, tag_list_array});
}
class MockIndexFilter : public IndexFilter {
public:
bool is_filtered(uint64_t id) const override {
return id % 2 == 1;
}
};
inline arrow::Result<std::shared_ptr<Table>> TakeRowsByIndices(
const std::shared_ptr<Table> &table, const std::vector<int> &row_indices) {
arrow::MemoryPool *pool = arrow::default_memory_pool();
arrow::Int32Builder indices_builder(pool);
ARROW_RETURN_NOT_OK(
indices_builder.AppendValues(row_indices.data(), row_indices.size()));
std::shared_ptr<arrow::Array> indices_array;
ARROW_RETURN_NOT_OK(indices_builder.Finish(&indices_array));
// 2. 对每一列执行 Take 操作
std::vector<std::shared_ptr<arrow::ChunkedArray>> new_columns;
for (const auto &column : table->columns()) {
// 使用 Take 提取指定索引的元素
ARROW_ASSIGN_OR_RAISE(auto taken_array, cp::Take(column, indices_array));
new_columns.emplace_back(taken_array.chunked_array());
}
// 3. 构造新的 Table
return arrow::Table::Make(table->schema(), new_columns, row_indices.size());
}
class MockSegment : public Segment {
public:
MockSegment() : Segment() {}
virtual ~MockSegment() = default;
SegmentID id() const override {
return 0;
}
TablePtr fetch(const std::vector<std::string> &columns,
const std::vector<int> &indices) const override {
std::string s = "";
for (auto i : indices) {
s += std::to_string(i);
s += ",";
}
LOG_INFO("Fetch indices: %s %s", get_column_names(columns).c_str(),
s.c_str());
auto table = CreateTable(1000).MoveValueUnsafe();
auto res = TakeRowsByIndices(table, indices);
if (!res.ok()) {
LOG_ERROR("Take error: %s", res.status().ToString().c_str());
return nullptr;
}
LOG_INFO("Take: %s", res.ValueOrDie()->ToString().c_str());
return res.MoveValueUnsafe();
}
ExecBatchPtr fetch(const std::vector<std::string> &columns,
int segment_doc_id) const override {
LOG_ERROR("Not implemented");
return nullptr;
}
static std::string get_column_names(const std::vector<std::string> &columns) {
std::string s = "";
for (auto i : columns) {
s += i;
s += ",";
}
return s;
}
RecordBatchReaderPtr scan(
const std::vector<std::string> &columns) const override {
auto table = CreateTable(10000);
LOG_INFO("Scan return: %s %s", get_column_names(columns).c_str(),
table.ValueOrDie()->ToString().c_str());
return std::make_shared<arrow::TableBatchReader>(table.ValueOrDie());
}
const IndexFilter::Ptr get_filter() override {
return std::make_shared<MockIndexFilter>();
}
CombinedVectorColumnIndexer::Ptr get_quant_combined_vector_indexer(
const std::string &field_name) const override {
return std::make_shared<MockVectorIndexer>();
}
CombinedVectorColumnIndexer::Ptr get_combined_vector_indexer(
const std::string &field_name) const override {
return std::make_shared<MockVectorIndexer>();
}
InvertedColumnIndexer::Ptr get_scalar_indexer(
const std::string &field_name) const override {
return std::make_shared<MockInvertIndexer>();
}
SegmentMeta::Ptr meta() const override {
return nullptr;
}
uint64_t doc_count(const IndexFilter::Ptr filter = nullptr) override {
return 0;
}
Status add_column(FieldSchema::Ptr column_schema,
const std::string &expression,
const AddColumnOptions &options) override {
return Status::InternalError();
}
Status alter_column(const std::string &column_name,
const FieldSchema::Ptr &new_column_schema,
const AlterColumnOptions &options) override {
return Status::InternalError();
}
Status drop_column(const std::string &column_name) override {
return Status::OK();
}
Status create_all_vector_index(
int concurrency, SegmentMeta::Ptr *new_segmnet_meta,
std::unordered_map<std::string, VectorColumnIndexer::Ptr>
*vector_indexers,
std::unordered_map<std::string, VectorColumnIndexer::Ptr>
*quant_vector_indexers) override {
return Status::OK();
}
Status create_vector_index(
const std::string &column, const IndexParams::Ptr &index_params,
int concurrency, SegmentMeta::Ptr *new_segmnet_meta,
std::unordered_map<std::string, VectorColumnIndexer::Ptr>
*vector_indexers,
std::unordered_map<std::string, VectorColumnIndexer::Ptr>
*quant_vector_indexers) override {
return Status::OK();
}
Status drop_vector_index(
const std::string &column, SegmentMeta::Ptr *new_segmnet_meta,
std::unordered_map<std::string, VectorColumnIndexer::Ptr>
*vector_indexers) override {
return Status::OK();
}
Status reload_vector_index(
const CollectionSchema &schema, const SegmentMeta::Ptr &segment_meta,
const std::unordered_map<std::string, VectorColumnIndexer::Ptr>
&vector_indexers,
const std::unordered_map<std::string, VectorColumnIndexer::Ptr>
&quant_vector_indexers) override {
return Status::OK();
}
bool vector_index_ready(const std::string &column,
const IndexParams::Ptr &index_params) const override {
return true;
}
bool all_vector_index_ready() const override {
return true;
}
Status create_scalar_index(
const std::vector<std::string> &columns,
const IndexParams::Ptr &index_params, SegmentMeta::Ptr *new_segment_meta,
InvertedIndexer::Ptr *new_scalar_indexer) override {
return Status::OK();
}
Status drop_scalar_index(const std::vector<std::string> &columns,
SegmentMeta::Ptr *new_segment_meta,
InvertedIndexer::Ptr *new_scalar_indexer) override {
return Status::OK();
}
Status reload_scalar_index(
const CollectionSchema &schema, const SegmentMeta::Ptr &segment_meta,
const InvertedIndexer::Ptr &scalar_indexer) override {
return Status::OK();
}
Status reload_fts_index(const CollectionSchema &schema,
const SegmentMeta::Ptr &segment_meta,
const FtsIndexer::Ptr &new_fts_indexer) override {
return Status::OK();
}
Status create_fts_index(const std::string &column,
const IndexParams::Ptr &index_params,
SegmentMeta::Ptr *new_segment_meta,
FtsIndexer::Ptr *output_fts_indexer) override {
return Status::OK();
}
Status drop_fts_index(const std::string &column,
SegmentMeta::Ptr *new_segment_meta,
FtsIndexer::Ptr *output_fts_indexer) override {
return Status::OK();
}
Status Insert(Doc &doc) override {
return Status::OK();
}
Status Upsert(Doc &doc) override {
return Status::OK();
}
Status Update(Doc &doc) override {
return Status::OK();
}
Status Delete(const std::string &pk) override {
return Status::OK();
}
Status Delete(uint64_t doc_id) override {
return Status::OK();
}
Doc::Ptr Fetch(uint64_t doc_id,
const std::optional<std::vector<std::string>> &output_fields =
std::nullopt,
bool include_vector = true) override {
return nullptr;
}
std::vector<VectorColumnIndexer::Ptr> get_vector_indexer(
const std::string &field_name) const override {
return {};
}
std::vector<VectorColumnIndexer::Ptr> get_quant_vector_indexer(
const std::string &field_name) const override {
return {};
}
fts::FtsColumnIndexerPtr get_fts_indexer(
const std::string &field_name) const override {
return nullptr;
}
Result<std::vector<fts::FtsResult>> fts_search(
const std::string &field_name, const fts::FtsAstNode &ast,
const fts::FtsQueryParams &params) override {
return std::vector<fts::FtsResult>{};
}
Status flush() override {
return Status::OK();
}
Status dump() override {
return Status::OK();
}
Status destroy() override {
return Status::OK();
}
};
} // namespace zvec
+322
View File
@@ -0,0 +1,322 @@
// 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 <gtest/gtest.h>
#include "db/sqlengine/analyzer/query_info_helper.h"
#include "db/sqlengine/sqlengine_impl.h"
#include "zvec/db/index_params.h"
// #define private public
#include <memory>
#include "db/sqlengine/planner/optimizer.h"
#include "mock_segment.h"
// #undef private
namespace zvec::sqlengine {
class MockInvertCondOptimizer : public InvertCondOptimizer {
public:
explicit MockInvertCondOptimizer(CollectionSchema *collection_schema)
: InvertCondOptimizer(collection_schema) {}
public:
bool invert_rule(Segment *segment, QueryRelNode *invert_cond) override;
};
bool MockInvertCondOptimizer::invert_rule(Segment *segment,
QueryRelNode *invert_cond) {
if (invert_cond->op() == QueryNodeOp::Q_IN) {
return true;
}
std::string invert_value = invert_cond->right()->text();
std::string numeric_text{""};
QueryInfoHelper::data_buf_2_text(invert_value, DataType::INT32,
&numeric_text);
int age = atoi(numeric_text.c_str());
std::cout << "invert cond: age is " << age << std::endl;
// invert cond as less than 100
if (age < 100) {
return true;
}
return false;
}
class OptimizerTest : public testing::Test {
public:
// Sets up the test fixture.
static void SetUpTestSuite() {
schema = std::make_shared<CollectionSchema>();
auto &collection_schema_ = *schema;
collection_schema_.set_name("collection");
// feature field
auto column1 = std::make_shared<FieldSchema>();
auto vector_params = std::make_shared<FlatIndexParams>(MetricType::IP);
column1->set_name("face_feature");
column1->set_index_params(vector_params);
column1->set_dimension(4);
column1->set_data_type(DataType::VECTOR_FP32);
collection_schema_.add_field(column1);
// invert field
auto column2 = std::make_shared<FieldSchema>();
column2->set_name("age");
column2->set_data_type(DataType::INT32);
column2->set_index_params(std::make_shared<InvertIndexParams>(false));
collection_schema_.add_field(column2);
}
// Tears down the test fixture.
static void TearDownTestSuite() {}
protected:
inline static CollectionSchema::Ptr schema;
Profiler::Ptr profiler_{new Profiler};
};
TEST_F(OptimizerTest, Basic) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.filter_ = "age > 200";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
Optimizer::Ptr optimizer =
std::make_shared<MockInvertCondOptimizer>(schema.get());
auto segment = std::make_shared<MockSegment>();
bool optimized = optimizer->optimize(segment.get(), query_info.get());
ASSERT_TRUE(optimized);
}
// case 1. invert subroot same as invert cond, do nothing
TEST_F(OptimizerTest, Case1) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.filter_ = "age > 12";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
Optimizer::Ptr optimizer =
std::make_shared<MockInvertCondOptimizer>(schema.get());
auto segment = std::make_shared<MockSegment>();
bool optimized = optimizer->optimize(segment.get(), query_info.get());
ASSERT_FALSE(optimized);
}
// case 2.1 invert subroot is not found, all conds are forward cond
TEST_F(OptimizerTest, Case2_1) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.filter_ = "age > 100 and age > 101 or age > 102";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
Optimizer::Ptr optimizer =
std::make_shared<MockInvertCondOptimizer>(schema.get());
auto segment = std::make_shared<MockSegment>();
bool optimized = optimizer->optimize(segment.get(), query_info.get());
ASSERT_TRUE(optimized);
}
// case 2.2 invert subroot is not found, some conds are forward cond
// while left invert cond cannot be invert cond any more
TEST_F(OptimizerTest, Case2_2) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.filter_ = "age > 100 or age > 90";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
Optimizer::Ptr optimizer =
std::make_shared<MockInvertCondOptimizer>(schema.get());
auto segment = std::make_shared<MockSegment>();
bool optimized = optimizer->optimize(segment.get(), query_info.get());
ASSERT_FALSE(optimized);
}
// case 3.1 subroot is found and be part of invert cond
TEST_F(OptimizerTest, Case3_1) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.filter_ = "age > 100 and age > 101 and age > 10";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
Optimizer::Ptr optimizer =
std::make_shared<MockInvertCondOptimizer>(schema.get());
auto segment = std::make_shared<MockSegment>();
bool optimized = optimizer->optimize(segment.get(), query_info.get());
ASSERT_TRUE(optimized);
ASSERT_TRUE(ret);
}
// case 3.2 subroot is found and be part of invert cond
TEST_F(OptimizerTest, Case3_2) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.filter_ = "age > 10 and age > 11 and age > 100";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
Optimizer::Ptr optimizer =
std::make_shared<MockInvertCondOptimizer>(schema.get());
auto segment = std::make_shared<MockSegment>();
bool optimized = optimizer->optimize(segment.get(), query_info.get());
ASSERT_TRUE(optimized);
}
// case 3.3 subroot is found and be part of invert cond
TEST_F(OptimizerTest, Case3_3) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.filter_ = "(age > 10 or age > 11) and age > 100";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
Optimizer::Ptr optimizer =
std::make_shared<MockInvertCondOptimizer>(schema.get());
auto segment = std::make_shared<MockSegment>();
bool optimized = optimizer->optimize(segment.get(), query_info.get());
ASSERT_TRUE(optimized);
}
// case 3.4 subroot is found and be part of invert cond, but others also have
// invert
TEST_F(OptimizerTest, Case3_4) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.filter_ = "age > 10 and (age > 101 and (age > 10 and age > 10))";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
Optimizer::Ptr optimizer =
std::make_shared<MockInvertCondOptimizer>(schema.get());
auto segment = std::make_shared<MockSegment>();
bool optimized = optimizer->optimize(segment.get(), query_info.get());
ASSERT_FALSE(optimized);
}
// case 4, optimize with in expr
TEST_F(OptimizerTest, Case4) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.filter_ = "age in (10, 20)";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr query_info = ret.value();
Optimizer::Ptr optimizer =
std::make_shared<MockInvertCondOptimizer>(schema.get());
auto segment = std::make_shared<MockSegment>();
bool optimized = optimizer->optimize(segment.get(), query_info.get());
// in will not optimized
ASSERT_FALSE(optimized);
// in and optimizable, optimize optimizable
query.filter_ = "age in (10, 20) and age > 100";
ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
query_info = ret.value();
optimized = optimizer->optimize(segment.get(), query_info.get());
ASSERT_TRUE(optimized);
// in or optimizable, not optimized
query.filter_ = "age in (10, 20) or age > 100";
ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
query_info = ret.value();
optimized = optimizer->optimize(segment.get(), query_info.get());
ASSERT_FALSE(optimized);
}
} // namespace zvec::sqlengine
+697
View File
@@ -0,0 +1,697 @@
// 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 <memory>
#include <gmock/gmock-matchers.h>
#include <gtest/gtest.h>
#include "db/sqlengine/sqlengine_impl.h"
#include "zvec/db/doc.h"
#include "zvec/db/schema.h"
#include "profiler.h"
namespace zvec::sqlengine {
class QueryInfoTest : public testing::Test {
public:
// Sets up the test fixture.
static void SetUpTestSuite() {
schema = std::make_shared<CollectionSchema>();
auto &param = *schema;
param.set_name("1collection");
auto column1 = std::make_shared<FieldSchema>();
auto vector_params = std::make_shared<FlatIndexParams>(MetricType::IP);
column1->set_name("face_feature");
column1->set_index_params(vector_params);
column1->set_dimension(4);
column1->set_data_type(DataType::VECTOR_FP32);
param.add_field(column1);
auto column2 = std::make_shared<FieldSchema>();
column2->set_name("name");
column2->set_data_type(DataType::UINT32);
param.add_field(column2);
auto column3 = std::make_shared<FieldSchema>();
column3->set_name("category");
column3->set_data_type(DataType::STRING);
param.add_field(column3);
auto column4 = std::make_shared<FieldSchema>();
column4->set_name("face_feature");
column4->set_dimension(4);
column4->set_data_type(DataType::VECTOR_FP32);
param.add_field(column4);
auto column5 = std::make_shared<FieldSchema>();
column5->set_name("1-dash_score_field");
column5->set_dimension(5);
column5->set_data_type(DataType::STRING);
param.add_field(column5);
{
auto column = std::make_shared<FieldSchema>();
column->set_name("name_array");
column->set_data_type(DataType::ARRAY_UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("category_array");
column->set_data_type(DataType::ARRAY_STRING);
param.add_field(column);
}
}
// Tears down the test fixture.
static void TearDownTestSuite() {}
protected:
Profiler::Ptr profiler_{new Profiler};
inline static CollectionSchema::Ptr schema;
};
TEST_F(QueryInfoTest, BasicQueryRequest) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.set_vector("[0.1, 0.2, 0.3, 0.4]");
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.target_.query_params_ = std::make_shared<QueryParams>(IndexType::FLAT);
query.target_.query_params_->set_radius(0.8F);
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
QueryInfo::Ptr new_query_info = ret.value();
auto &query_fields = new_query_info->query_fields();
EXPECT_EQ(query_fields.size(), 5);
EXPECT_EQ(query_fields[0]->field_name(), "name");
EXPECT_EQ(query_fields[1]->field_name(), "category");
EXPECT_EQ(query_fields[2]->field_name(), "1-dash_score_field");
EXPECT_EQ(query_fields[3]->field_name(), "name_array");
EXPECT_EQ(query_fields[4]->field_name(), "category_array");
EXPECT_EQ(new_query_info->query_topn(), 11);
EXPECT_FALSE(new_query_info->filter_cond());
EXPECT_FALSE(new_query_info->invert_cond());
EXPECT_FALSE(new_query_info->post_filter_cond());
EXPECT_FALSE(new_query_info->post_invert_cond());
ASSERT_TRUE(new_query_info->vector_cond_info());
auto vector_cond = new_query_info->vector_cond_info();
EXPECT_EQ(1, vector_cond->batch());
EXPECT_EQ("face_feature", vector_cond->vector_field_name());
EXPECT_EQ(query.target_.get_vector_clause()->query_vector_,
vector_cond->vector_term());
EXPECT_EQ(query.target_.get_vector_clause()->sparse_indices_,
vector_cond->vector_sparse_indices());
EXPECT_EQ(query.target_.get_vector_clause()->sparse_values_,
vector_cond->vector_sparse_values());
EXPECT_EQ(query.target_.query_params_, vector_cond->query_params());
}
TEST_F(QueryInfoTest, QueryRequestWithFilter) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.set_vector("[0.1, 0.2, 0.3, 0.4]");
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.target_.query_params_ = std::make_shared<QueryParams>(IndexType::FLAT);
query.target_.query_params_->set_radius(0.8F);
query.filter_ = "name<3 or name=4 or 1-dash_score_field='test'";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
auto &query_fields = new_query_info->query_fields();
EXPECT_EQ(query_fields.size(), 5);
EXPECT_EQ(query_fields[0]->field_name(), "name");
EXPECT_EQ(query_fields[1]->field_name(), "category");
EXPECT_EQ(query_fields[2]->field_name(), "1-dash_score_field");
EXPECT_EQ(query_fields[3]->field_name(), "name_array");
EXPECT_EQ(query_fields[4]->field_name(), "category_array");
EXPECT_EQ(new_query_info->query_topn(), 11);
EXPECT_TRUE(new_query_info->filter_cond());
EXPECT_FALSE(new_query_info->invert_cond());
EXPECT_FALSE(new_query_info->post_filter_cond());
EXPECT_FALSE(new_query_info->post_invert_cond());
ASSERT_TRUE(new_query_info->vector_cond_info());
auto vector_cond = new_query_info->vector_cond_info();
EXPECT_EQ(1, vector_cond->batch());
EXPECT_EQ("face_feature", vector_cond->vector_field_name());
EXPECT_EQ(query.target_.get_vector_clause()->query_vector_,
vector_cond->vector_term());
EXPECT_EQ(query.target_.get_vector_clause()->sparse_indices_,
vector_cond->vector_sparse_indices());
EXPECT_EQ(query.target_.get_vector_clause()->sparse_values_,
vector_cond->vector_sparse_values());
EXPECT_EQ(query.target_.query_params_, vector_cond->query_params());
EXPECT_TRUE(new_query_info->filter_cond());
// (nullptr) and (xxx)
auto filter_cond = new_query_info->filter_cond();
EXPECT_EQ(filter_cond->op_name(), "and");
EXPECT_FALSE(filter_cond->left());
// ((name<3) or (name=4)) or (1-dash_score_field=test)
auto right = std::dynamic_pointer_cast<QueryNode>(filter_cond->right());
EXPECT_TRUE(right);
EXPECT_EQ(right->op_name(), "or");
// 1-dash_score_field=test
auto number_field_filter =
std::dynamic_pointer_cast<QueryNode>(right->right());
ASSERT_TRUE(number_field_filter);
EXPECT_EQ(number_field_filter->op_name(), "=");
auto left_key =
std::dynamic_pointer_cast<QueryIDNode>(number_field_filter->left());
EXPECT_EQ(left_key->op_name(), "ID");
EXPECT_EQ(left_key->value(), "1-dash_score_field");
auto right_const = std::dynamic_pointer_cast<QueryConstantNode>(
number_field_filter->right());
ASSERT_TRUE(right_const);
EXPECT_EQ(right_const->op_name(), "STRING_VALUE");
EXPECT_EQ(right_const->value(), "test");
// (name<3) or (name=4)
auto left = std::dynamic_pointer_cast<QueryNode>(right->left());
ASSERT_TRUE(left);
EXPECT_EQ(left->op_name(), "or");
auto or1 = std::dynamic_pointer_cast<QueryNode>(left->left());
EXPECT_EQ(or1->op_name(), "<");
auto id1 = std::dynamic_pointer_cast<QueryIDNode>(or1->left());
ASSERT_TRUE(id1);
EXPECT_EQ(id1->op_name(), "ID");
EXPECT_EQ(id1->value(), "name");
auto const1 = std::dynamic_pointer_cast<QueryConstantNode>(or1->right());
ASSERT_TRUE(const1);
EXPECT_EQ(const1->op_name(), "INT_VALUE");
EXPECT_EQ(const1->value(), "3");
auto or2 = std::dynamic_pointer_cast<QueryNode>(left->right());
EXPECT_EQ(or2->op_name(), "=");
}
TEST_F(QueryInfoTest, QueryRequestWithIncludeVector) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.set_vector("[0.1, 0.2, 0.3, 0.4]");
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.target_.query_params_ = std::make_shared<QueryParams>(IndexType::FLAT);
query.target_.query_params_->set_radius(0.8F);
query.include_vector_ = true;
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
auto &query_fields = new_query_info->query_fields();
EXPECT_EQ(query_fields.size(), 6);
EXPECT_EQ(query_fields[0]->field_name(), "name");
EXPECT_EQ(query_fields[1]->field_name(), "category");
EXPECT_EQ(query_fields[2]->field_name(), "1-dash_score_field");
EXPECT_EQ(query_fields[3]->field_name(), "name_array");
EXPECT_EQ(query_fields[4]->field_name(), "category_array");
EXPECT_EQ(query_fields[5]->field_name(), "face_feature");
EXPECT_EQ(new_query_info->query_topn(), 11);
EXPECT_FALSE(new_query_info->filter_cond());
EXPECT_FALSE(new_query_info->invert_cond());
EXPECT_FALSE(new_query_info->post_filter_cond());
EXPECT_FALSE(new_query_info->post_invert_cond());
ASSERT_TRUE(new_query_info->vector_cond_info());
auto vector_cond = new_query_info->vector_cond_info();
EXPECT_EQ(1, vector_cond->batch());
EXPECT_EQ("face_feature", vector_cond->vector_field_name());
EXPECT_EQ(query.target_.get_vector_clause()->query_vector_,
vector_cond->vector_term());
EXPECT_EQ(query.target_.get_vector_clause()->sparse_indices_,
vector_cond->vector_sparse_indices());
EXPECT_EQ(query.target_.get_vector_clause()->sparse_values_,
vector_cond->vector_sparse_values());
EXPECT_EQ(query.target_.query_params_, vector_cond->query_params());
}
TEST_F(QueryInfoTest, OR_ANCESTOR) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.set_vector("[0.1, 0.2, 0.3, 0.4]");
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.target_.query_params_ = std::make_shared<QueryParams>(IndexType::FLAT);
query.target_.query_params_->set_radius(0.8F);
query.filter_ = "name=1 and (name=2 or name=3)";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
}
TEST_F(QueryInfoTest, QueryRequestWithInFilter) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 10;
query.target_.set_vector("[0.1, 0.2, 0.3, 0.4]");
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.target_.query_params_ = std::make_shared<QueryParams>(IndexType::FLAT);
query.target_.query_params_->set_radius(0.8F);
query.filter_ =
"name=3 or name in (1, 2, 3) or category not in (\"a\", \"b\", \"c\")";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
auto &query_fields = new_query_info->query_fields();
EXPECT_EQ(query_fields.size(), 5);
EXPECT_EQ(query_fields[0]->field_name(), "name");
EXPECT_EQ(query_fields[1]->field_name(), "category");
EXPECT_EQ(query_fields[2]->field_name(), "1-dash_score_field");
EXPECT_EQ(query_fields[3]->field_name(), "name_array");
EXPECT_EQ(query_fields[4]->field_name(), "category_array");
EXPECT_EQ(new_query_info->query_topn(), 10);
EXPECT_FALSE(new_query_info->invert_cond());
EXPECT_FALSE(new_query_info->post_filter_cond());
EXPECT_FALSE(new_query_info->post_invert_cond());
ASSERT_TRUE(new_query_info->vector_cond_info());
auto vector_cond = new_query_info->vector_cond_info();
EXPECT_EQ(1, vector_cond->batch());
EXPECT_EQ("face_feature", vector_cond->vector_field_name());
std::vector<float> data{1.1, 2.2, 3.3, 4.4};
EXPECT_EQ(query.target_.get_vector_clause()->query_vector_,
vector_cond->vector_term());
EXPECT_TRUE(new_query_info->filter_cond());
// (nullptr) and (xxx)
auto filter_cond = new_query_info->filter_cond();
EXPECT_EQ(filter_cond->op_name(), "and");
EXPECT_FALSE(filter_cond->left());
// ((name=3) or (name in (1, 2, 3))) or (category not in ("a", "b", "c"))
auto right = std::dynamic_pointer_cast<QueryNode>(filter_cond->right());
EXPECT_TRUE(right);
EXPECT_EQ(right->op_name(), "or");
// category in ("a", "b", "c")
auto category_filter = std::dynamic_pointer_cast<QueryNode>(right->right());
ASSERT_TRUE(category_filter);
EXPECT_EQ(category_filter->op_name(), " in ");
auto left_key =
std::dynamic_pointer_cast<QueryIDNode>(category_filter->left());
EXPECT_EQ(left_key->op_name(), "ID");
EXPECT_EQ(left_key->value(), "category");
auto right_const =
std::dynamic_pointer_cast<QueryListNode>(category_filter->right());
ASSERT_TRUE(right_const);
EXPECT_EQ(right_const->op_name(), "LIST_VALUE");
EXPECT_EQ(right_const->text(), "NOT (a, b, c)");
// (name=3) or (name in (1, 2, 3))
auto left = std::dynamic_pointer_cast<QueryNode>(right->left());
ASSERT_TRUE(left);
EXPECT_EQ(left->op_name(), "or");
auto or1 = std::dynamic_pointer_cast<QueryNode>(left->left());
EXPECT_EQ(or1->op_name(), "=");
auto id1 = std::dynamic_pointer_cast<QueryIDNode>(or1->left());
ASSERT_TRUE(id1);
EXPECT_EQ(id1->op_name(), "ID");
EXPECT_EQ(id1->value(), "name");
auto const1 = std::dynamic_pointer_cast<QueryConstantNode>(or1->right());
ASSERT_TRUE(const1);
EXPECT_EQ(const1->op_name(), "INT_VALUE");
EXPECT_EQ(const1->value(), "3");
auto or2 = std::dynamic_pointer_cast<QueryNode>(left->right());
EXPECT_EQ(or2->op_name(), " in ");
auto id2 = std::dynamic_pointer_cast<QueryIDNode>(or2->left());
ASSERT_TRUE(id2);
EXPECT_EQ(id2->op_name(), "ID");
EXPECT_EQ(id2->value(), "name");
auto const2 = std::dynamic_pointer_cast<QueryListNode>(or2->right());
ASSERT_TRUE(const2);
EXPECT_EQ(const2->op_name(), "LIST_VALUE");
EXPECT_EQ(const2->text(), "(1, 2, 3)");
}
TEST_F(QueryInfoTest, QueryRequestWithInFilterWrong) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.target_.set_vector("[0.1, 0.2, 0.3, 0.4]");
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.target_.query_params_ = std::make_shared<QueryParams>(IndexType::FLAT);
query.target_.query_params_->set_radius(0.8F);
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
query.filter_ = ("name in ()");
ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
query.filter_ = ("name in (\"a\", 2, 3)");
ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
query.filter_ = ("name in (1.1, 2, 3)");
ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
query.filter_ = ("category in (1.1, \"b\")");
ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
}
TEST_F(QueryInfoTest, QueryRequestWithInFilterNum1024) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 10;
query.target_.set_vector("[0.1, 0.2, 0.3, 0.4]");
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.target_.query_params_ = std::make_shared<QueryParams>(IndexType::FLAT);
query.target_.query_params_->set_radius(0.8F);
std::string filter_str;
for (int i = 0; i < 1024; i++) {
if (i != 0) {
filter_str += " or ";
}
filter_str += "name=" + std::to_string(i);
}
query.filter_ = filter_str;
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
auto &query_fields = new_query_info->query_fields();
EXPECT_EQ(query_fields.size(), 5);
EXPECT_EQ(query_fields[0]->field_name(), "name");
EXPECT_EQ(query_fields[1]->field_name(), "category");
EXPECT_EQ(query_fields[2]->field_name(), "1-dash_score_field");
EXPECT_EQ(query_fields[3]->field_name(), "name_array");
EXPECT_EQ(query_fields[4]->field_name(), "category_array");
EXPECT_EQ(new_query_info->query_topn(), 10);
EXPECT_FALSE(new_query_info->invert_cond());
EXPECT_FALSE(new_query_info->post_filter_cond());
EXPECT_FALSE(new_query_info->post_invert_cond());
ASSERT_TRUE(new_query_info->vector_cond_info());
auto vector_cond = new_query_info->vector_cond_info();
EXPECT_EQ(1, vector_cond->batch());
EXPECT_EQ("face_feature", vector_cond->vector_field_name());
}
TEST_F(QueryInfoTest, QueryRequestWithFilter_contain) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 10;
query.target_.set_vector("[0.1, 0.2, 0.3, 0.4]");
query.target_.field_name_ = "face_feature";
query.include_vector_ = false;
query.target_.query_params_ = std::make_shared<QueryParams>(IndexType::FLAT);
query.target_.query_params_->set_radius(0.8F);
query.filter_ =
R"( name_array contain_all (1, 2, 3) and )"
R"( (name_array not contain_all (4, 5) or category_array contain_any
("a", "b")) )"
R"( or category_array not contain_any ("c", "d", "e")
)";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
auto &query_fields = new_query_info->query_fields();
// pre-defined schema field
EXPECT_EQ(query_fields.size(), 5);
EXPECT_EQ(query_fields[0]->field_name(), "name");
EXPECT_EQ(query_fields[1]->field_name(), "category");
EXPECT_EQ(query_fields[2]->field_name(), "1-dash_score_field");
EXPECT_EQ(query_fields[3]->field_name(), "name_array");
EXPECT_EQ(query_fields[4]->field_name(), "category_array");
EXPECT_EQ(new_query_info->query_topn(), 10);
EXPECT_FALSE(new_query_info->invert_cond());
EXPECT_FALSE(new_query_info->post_filter_cond());
EXPECT_FALSE(new_query_info->post_invert_cond());
ASSERT_TRUE(new_query_info->vector_cond_info());
auto vector_cond = new_query_info->vector_cond_info();
EXPECT_EQ(1, vector_cond->batch());
EXPECT_EQ("face_feature", vector_cond->vector_field_name());
EXPECT_TRUE(new_query_info->filter_cond());
/*
_________________[and]__________________
/ \
[nullptr(vector_cond)] [filter condition]
*/
// (nullptr) and (xxx)
auto filter_cond = new_query_info->filter_cond();
EXPECT_EQ(filter_cond->op_name(), "and");
EXPECT_FALSE(filter_cond->left());
/*
_______________[or]_______________
/ \
_____________[and]_____________ [category_array not
contain_any
("c", "d", "e")]
/ \
[name_array contain_all (1, 2, 3)] ___________[or]______________
/ \
[name_array not contain_all (4, 5)] [category_array
contain_any ("a", "b")]
*/
// name_array contain_all (1, 2, 3) and
// (name_array not contain_all (4, 5) or category_array contain_any ("a",
// "b")) or category_array not contain_any ("c", "d", "e")
auto parent_node = std::dynamic_pointer_cast<QueryNode>(filter_cond);
auto cur_node = std::dynamic_pointer_cast<QueryNode>(filter_cond->right());
EXPECT_TRUE(cur_node);
EXPECT_EQ(cur_node->op_name(), "or");
// category_array not contain_any ("c", "d", "e")
parent_node = std::dynamic_pointer_cast<QueryNode>(cur_node);
cur_node = std::dynamic_pointer_cast<QueryNode>(cur_node->right());
EXPECT_TRUE(cur_node);
EXPECT_EQ(cur_node->op_name(), " contain_any ");
{
auto left_key = std::dynamic_pointer_cast<QueryIDNode>(cur_node->left());
EXPECT_EQ(left_key->op_name(), "ID");
EXPECT_EQ(left_key->value(), "category_array");
auto right_const =
std::dynamic_pointer_cast<QueryListNode>(cur_node->right());
ASSERT_TRUE(right_const);
EXPECT_EQ(right_const->op_name(), "LIST_VALUE");
EXPECT_EQ(right_const->text(), "NOT (c, d, e)");
}
cur_node = parent_node;
// name_array contain_all (1, 2, 3) and
// (name_array not contain_all (4, 5) or category_array contain_any ("a",
// "b"))
parent_node = std::dynamic_pointer_cast<QueryNode>(cur_node);
cur_node = std::dynamic_pointer_cast<QueryNode>(cur_node->left());
EXPECT_TRUE(cur_node);
EXPECT_EQ(cur_node->op_name(), "and");
// the left side of 'and'
// name_array contain_all (1, 2, 3)
parent_node = std::dynamic_pointer_cast<QueryNode>(cur_node);
cur_node = std::dynamic_pointer_cast<QueryNode>(cur_node->left());
EXPECT_TRUE(cur_node);
EXPECT_EQ(cur_node->op_name(), " contain_all ");
{
auto left_key = std::dynamic_pointer_cast<QueryIDNode>(cur_node->left());
EXPECT_EQ(left_key->op_name(), "ID");
EXPECT_EQ(left_key->value(), "name_array");
auto right_const =
std::dynamic_pointer_cast<QueryListNode>(cur_node->right());
ASSERT_TRUE(right_const);
EXPECT_EQ(right_const->op_name(), "LIST_VALUE");
EXPECT_EQ(right_const->text(), "(1, 2, 3)");
}
cur_node = parent_node;
// the right side of 'and'
// (name_array not contain_all (4, 5) or category_array contain_any ("a",
// "b"))
parent_node = std::dynamic_pointer_cast<QueryNode>(cur_node);
cur_node = std::dynamic_pointer_cast<QueryNode>(cur_node->right());
EXPECT_TRUE(cur_node);
EXPECT_EQ(cur_node->op_name(), "or");
// name_array not contain_all (4, 5)
parent_node = std::dynamic_pointer_cast<QueryNode>(cur_node);
cur_node = std::dynamic_pointer_cast<QueryNode>(cur_node->left());
EXPECT_TRUE(cur_node);
EXPECT_EQ(cur_node->op_name(), " contain_all ");
{
auto left_key = std::dynamic_pointer_cast<QueryIDNode>(cur_node->left());
EXPECT_EQ(left_key->op_name(), "ID");
EXPECT_EQ(left_key->value(), "name_array");
auto right_const =
std::dynamic_pointer_cast<QueryListNode>(cur_node->right());
ASSERT_TRUE(right_const);
EXPECT_EQ(right_const->op_name(), "LIST_VALUE");
EXPECT_EQ(right_const->text(), "NOT (4, 5)");
}
cur_node = parent_node;
// category_array contain_any ("a", "b"))
parent_node = std::dynamic_pointer_cast<QueryNode>(cur_node);
cur_node = std::dynamic_pointer_cast<QueryNode>(cur_node->right());
EXPECT_TRUE(cur_node);
EXPECT_EQ(cur_node->op_name(), " contain_any ");
{
auto left_key = std::dynamic_pointer_cast<QueryIDNode>(cur_node->left());
EXPECT_EQ(left_key->op_name(), "ID");
EXPECT_EQ(left_key->value(), "category_array");
auto right_const =
std::dynamic_pointer_cast<QueryListNode>(cur_node->right());
ASSERT_TRUE(right_const);
EXPECT_EQ(right_const->op_name(), "LIST_VALUE");
EXPECT_EQ(right_const->text(), "(a, b)");
}
cur_node = parent_node;
}
TEST_F(QueryInfoTest, SelectNonExistField) {
SearchQuery query;
query.output_fields_ = {"category_array", "not_exist_field"};
query.topk_ = 11;
query.include_vector_ = false;
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(ret.error().message(),
testing::HasSubstr("not defined in schema"));
}
TEST_F(QueryInfoTest, ContainAllExceedLimit) {
SearchQuery query;
query.topk_ = 200;
query.filter_ = "name_array not contain_all (";
for (int i = 0; i <= 32; i++) {
query.filter_ += std::to_string(i);
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(ret.error().message(),
testing::HasSubstr(
"Contain_* rel expr only support list size no more than 32"));
}
TEST_F(QueryInfoTest, ContainAnyExceedLimit) {
SearchQuery query;
query.topk_ = 200;
query.filter_ = "name_array not contain_any (";
for (int i = 0; i <= 32; i++) {
query.filter_ += std::to_string(i);
if (i < 32) {
query.filter_ += ", ";
}
}
query.filter_ += ")";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(ret.error().message(),
testing::HasSubstr(
"Contain_* rel expr only support list size no more than 32"));
}
TEST_F(QueryInfoTest, ArrayLengthNonExistField) {
SearchQuery query;
query.topk_ = 200;
query.filter_ = "array_length(not_exist_field) > 1";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(ret.error().message(),
testing::HasSubstr("array_length argument not found in schema"));
}
TEST_F(QueryInfoTest, ArrayLengthOnNonArrayField) {
SearchQuery query;
query.topk_ = 200;
query.filter_ = "array_length(name) > 1";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(ret.error().message(),
testing::HasSubstr("array_length only support array"));
}
TEST_F(QueryInfoTest, ArrayLengthInvalidArgument) {
SearchQuery query;
query.topk_ = 200;
query.filter_ = "array_length(name_array) > '1'";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(
ret.error().message(),
testing::HasSubstr("array_length right side only support integer"));
}
TEST_F(QueryInfoTest, ArrayLengthInvalidOp) {
SearchQuery query;
query.topk_ = 200;
query.filter_ = "array_length(name_array) like '%'";
auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
auto ret = engine->build_query_info(schema, query, nullptr);
ASSERT_FALSE(ret.has_value());
EXPECT_THAT(ret.error().message(), testing::HasSubstr("syntax error"));
}
} // namespace zvec::sqlengine
+270
View File
@@ -0,0 +1,270 @@
// 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
#pragma once
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <memory>
#include <arrow/api.h>
#include <arrow/io/api.h>
#include <arrow/ipc/api.h>
#include <gtest/gtest.h>
#include "db/common/file_helper.h"
#include "db/index/common/version_manager.h"
#include "db/index/segment/segment.h"
#include "zvec/db/index_params.h"
#include "zvec/db/schema.h"
#include "zvec/db/type.h"
namespace zvec {
inline CollectionSchema::Ptr GetCollectionSchema() {
auto invert_params = std::make_shared<InvertIndexParams>(true);
auto collection_schema = std::make_shared<CollectionSchema>(
"test_collection",
std::vector<FieldSchema::Ptr>{
std::make_shared<FieldSchema>("id", DataType::UINT64, false, nullptr),
std::make_shared<FieldSchema>("invert_id", DataType::UINT64, false,
invert_params),
std::make_shared<FieldSchema>("bool", DataType::BOOL, false, nullptr),
std::make_shared<FieldSchema>("invert_bool", DataType::BOOL, false,
invert_params),
std::make_shared<FieldSchema>("bool_array", DataType::ARRAY_BOOL,
false, nullptr),
std::make_shared<FieldSchema>(
"invert_bool_array", DataType::ARRAY_BOOL, false, invert_params),
std::make_shared<FieldSchema>("name", DataType::STRING, false,
nullptr),
std::make_shared<FieldSchema>("invert_name", DataType::STRING, false,
invert_params),
std::make_shared<FieldSchema>("age", DataType::INT32, false, nullptr),
std::make_shared<FieldSchema>(
"invert_age", DataType::INT32, false,
std::make_shared<InvertIndexParams>(true)),
std::make_shared<FieldSchema>("score", DataType::DOUBLE, false,
nullptr),
std::make_shared<FieldSchema>("optional_age", DataType::UINT32, true,
nullptr),
std::make_shared<FieldSchema>("invert_optional_age", DataType::UINT32,
true, invert_params),
std::make_shared<FieldSchema>("category_set", DataType::ARRAY_INT32,
true, nullptr),
std::make_shared<FieldSchema>("invert_category_set",
DataType::ARRAY_INT32, true,
invert_params),
// add vector field
std::make_shared<FieldSchema>(
"dense", DataType::VECTOR_FP32, 4, false,
std::make_shared<FlatIndexParams>(MetricType::L2)),
// add sparse vector
std::make_shared<FieldSchema>(
"sparse", DataType::SPARSE_VECTOR_FP32, 0, false,
std::make_shared<FlatIndexParams>(MetricType::IP)),
});
return collection_schema;
}
inline Doc CreateDoc(const uint64_t doc_id) {
Doc new_doc;
new_doc.set_pk("pk_" + std::to_string(doc_id));
new_doc.set_doc_id(doc_id);
new_doc.set<uint64_t>("id", doc_id);
new_doc.set<uint64_t>("invert_id", doc_id);
new_doc.set<bool>("bool", doc_id % 100 == 0);
new_doc.set<bool>("invert_bool", doc_id % 100 == 0);
new_doc.set<int32_t>("age", doc_id % 100);
new_doc.set<int32_t>("invert_age", doc_id % 100);
if (uint32_t v = doc_id % 100; v) {
new_doc.set("optional_age", v);
new_doc.set("invert_optional_age", v);
}
auto name = "user_" + std::to_string(doc_id % 100);
new_doc.set<std::string>("name", name);
new_doc.set<std::string>("invert_name", name);
new_doc.set<double>("score", static_cast<double>(rand() % 1000) / 10.0);
// vector
std::vector<float> vv;
for (uint32_t i = 0; i < 4; i++) {
vv.push_back(static_cast<float>(doc_id));
}
new_doc.set<std::vector<float>>("dense", vv);
// sparse vector
{
std::vector<uint32_t> indices;
std::vector<float> values;
for (uint32_t i = 0; i < doc_id % 100; i++) {
indices.push_back(i);
values.push_back(static_cast<float>(doc_id));
}
new_doc.set<std::pair<std::vector<uint32_t>, std::vector<float>>>(
"sparse", std::make_pair(indices, values));
}
auto category_size = doc_id % 100;
if (category_size > 0) {
std::vector<int32_t> category;
for (uint32_t i = 1; i <= category_size; i++) {
category.push_back(i);
}
new_doc.set("category_set", category);
new_doc.set("invert_category_set", category);
}
if (doc_id % 3 == 0) {
new_doc.set<std::vector<bool>>("bool_array", {true, false, true});
new_doc.set<std::vector<bool>>("invert_bool_array", {true, false, true});
} else if (doc_id % 3 == 1) {
new_doc.set<std::vector<bool>>("bool_array", {true, true, true});
new_doc.set<std::vector<bool>>("invert_bool_array", {true, true, true});
} else {
new_doc.set<std::vector<bool>>("bool_array", {false, false, false});
new_doc.set<std::vector<bool>>("invert_bool_array", {false, false, false});
}
return new_doc;
}
inline Status InsertDoc(const Segment::Ptr &segment,
const uint64_t start_doc_id,
const uint64_t end_doc_id) {
srand(time(nullptr));
long long create_total = 0;
long long insert_total = 0;
for (auto doc_id = start_doc_id; doc_id < end_doc_id; doc_id++) {
if (segment) {
auto start = std::chrono::system_clock::now();
Doc new_doc = CreateDoc(doc_id);
auto end = std::chrono::system_clock::now();
auto create_cost =
std::chrono::duration_cast<std::chrono::microseconds>(end - start)
.count();
create_total += create_cost;
start = std::chrono::system_clock::now();
auto status = segment->Insert(new_doc);
if (!status.ok()) {
return status;
}
end = std::chrono::system_clock::now();
auto insert_cost =
std::chrono::duration_cast<std::chrono::microseconds>(end - start)
.count();
insert_total += insert_cost;
}
}
std::cout << "pure create cost " << create_total << "us" << std::endl;
std::cout << "pure insert cost " << insert_total << "us" << std::endl;
return Status::OK();
}
class RecallTest : public testing::Test {
protected:
static void SetUpTestSuite() {
FileHelper::RemoveDirectory(seg_path_);
FileHelper::CreateDirectory(seg_path_);
collection_schema_ = GetCollectionSchema();
auto segment = create_segment();
if (segment == nullptr) {
LOG_ERROR("create segment failed");
EXPECT_TRUE(segment != nullptr);
std::exit(EXIT_FAILURE);
}
auto status = InsertDoc(segment, 0, 10000);
if (!status.ok()) {
LOG_ERROR("insert doc failed: %s", status.c_str());
EXPECT_TRUE(status.ok());
std::exit(EXIT_FAILURE);
}
segments_.push_back(segment);
}
static void TearDownTestSuite() {
segments_.clear();
FileHelper::RemoveDirectory(seg_path_);
}
public:
static std::string GetPath() {
return seg_path_;
}
static Segment::Ptr create_segment();
protected:
static inline std::string seg_path_ = "./test_collection";
static inline CollectionSchema::Ptr collection_schema_;
static inline std::vector<Segment::Ptr> segments_;
};
inline Segment::Ptr RecallTest::create_segment() {
auto seg_path = GetPath();
auto segment_meta = std::make_shared<SegmentMeta>();
segment_meta->set_id(0);
auto id_map = IDMap::CreateAndOpen("test_collection", GetPath() + "/id_map",
true, false);
auto delete_store = std::make_shared<DeleteStore>("test_collection");
Version v1;
v1.set_schema(*collection_schema_);
std::string v_path = GetPath() + "/test_manifest";
FileHelper::CreateDirectory(v_path);
auto vm = VersionManager::Create(v_path, v1);
if (!vm.has_value()) {
LOG_ERROR("create version manager failed: %s", vm.error().c_str());
return nullptr;
}
BlockMeta mem_block;
mem_block.id_ = 0;
mem_block.type_ = BlockType::SCALAR;
mem_block.min_doc_id_ = 0;
mem_block.max_doc_id_ = 0;
mem_block.doc_count_ = 0;
segment_meta->set_writing_forward_block(mem_block);
SegmentOptions options;
options.read_only_ = false;
options.enable_mmap_ = true;
options.max_buffer_size_ = 256 * 1024;
auto result =
Segment::CreateAndOpen(GetPath(), *collection_schema_, 0, 0, id_map,
delete_store, vm.value(), options);
if (!result) {
LOG_ERROR("create segment failed: %s", result.error().c_str());
return nullptr;
}
auto segment = result.value();
return segment;
}
} // namespace zvec
+650
View File
@@ -0,0 +1,650 @@
// 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 "sqlengine/analyzer/simple_rewriter.h"
#include <gtest/gtest.h>
#include "db/sqlengine/analyzer/query_info.h"
#include "db/sqlengine/sqlengine_impl.h"
#include "zvec/db/doc.h"
#include "zvec/db/schema.h"
namespace zvec::sqlengine {
class SimpleRewriterTest : public testing::Test {
public:
// Sets up the test fixture.
static void SetUpTestSuite() {
schema = std::make_shared<CollectionSchema>();
auto &param = *schema;
param.set_name("1collection");
auto column1 = std::make_shared<FieldSchema>();
auto vector_params = std::make_shared<FlatIndexParams>(MetricType::IP);
column1->set_name("face_feature");
column1->set_index_params(vector_params);
column1->set_dimension(4);
column1->set_data_type(DataType::VECTOR_FP32);
param.add_field(column1);
auto column2 = std::make_shared<FieldSchema>();
column2->set_name("age");
column2->set_data_type(DataType::UINT32);
param.add_field(column2);
auto column_gender = std::make_shared<FieldSchema>();
column_gender->set_name("gender");
column_gender->set_data_type(DataType::UINT32);
param.add_field(column_gender);
auto column3 = std::make_shared<FieldSchema>();
column3->set_name("category");
column3->set_data_type(DataType::STRING);
param.add_field(column3);
auto column4 = std::make_shared<FieldSchema>();
column4->set_name("face_feature");
column4->set_dimension(4);
column4->set_data_type(DataType::VECTOR_FP32);
param.add_field(column4);
auto column5 = std::make_shared<FieldSchema>();
column5->set_name("filename");
column5->set_dimension(5);
column5->set_data_type(DataType::STRING);
param.add_field(column5);
{
auto column = std::make_shared<FieldSchema>();
column->set_name("loc");
column->set_data_type(DataType::UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("fid");
column->set_data_type(DataType::UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("agent_id");
column->set_data_type(DataType::UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("state");
column->set_data_type(DataType::UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("categoryId");
column->set_data_type(DataType::UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("passed_days");
column->set_data_type(DataType::UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("category_in");
column->set_data_type(DataType::UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("category_out");
column->set_data_type(DataType::UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("intAttr");
column->set_data_type(DataType::UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("intAttr");
column->set_data_type(DataType::UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("strAttr");
column->set_data_type(DataType::STRING);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("partitionName");
column->set_data_type(DataType::STRING);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("doc_id");
column->set_data_type(DataType::UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("a");
column->set_data_type(DataType::UINT32);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("is_type1");
column->set_data_type(DataType::BOOL);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("is_type2");
column->set_data_type(DataType::BOOL);
param.add_field(column);
}
{
auto column = std::make_shared<FieldSchema>();
column->set_name("category_array");
column->set_data_type(DataType::ARRAY_STRING);
param.add_field(column);
}
}
// Tears down the test fixture.
static void TearDownTestSuite() {}
QueryInfo::Ptr parse(const std::string &filter) {
SearchQuery query;
query.output_fields_ = {"*"};
query.topk_ = 11;
query.include_vector_ = false;
query.filter_ = filter;
auto engine = std::make_shared<SQLEngineImpl>(profiler_);
auto ret = engine->build_query_info(schema, query, nullptr);
// ASSERT_TRUE(ret.has_value());
QueryInfo::Ptr new_query_info = ret.value();
return new_query_info;
}
protected:
Profiler::Ptr profiler_{new Profiler};
inline static CollectionSchema::Ptr schema;
};
class EqOrRewriteTest : public SimpleRewriterTest {};
TEST_F(EqOrRewriteTest, SimpleEqOr) {
auto info = parse("age = 10 or age = 20 ");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(), "age in (10, 20)(FORWARD)");
}
TEST_F(EqOrRewriteTest, SimpleManyEqOr) {
auto info = parse(
"age = 1 or age = 2 or age = 3 or age = 4 "
"or age = 5 or age = 6 or age = 7 or age = 8 or age = 9 or age = 10 or "
"age = 11 or age = 12 or age = 13 or age = 14 or age = 15 or age = 16 or "
"age = 17 or age = 18 or age = 19 or age = 20 or age = 21 or age = 22 or "
"age = 23 or age = 24 or age = 25 or age = 26 or age = 27 or age = 28 or "
"age = 29 or age = 30 or age = 31 or age = 32 or age = 33 or age = 34 or "
"age = 35 or age = 36 or age = 37 or age = 38 or age = 39 or age = 40 or "
"age = 41 or age = 42 or age = 43 or age = 44 or age = 45 or age = 46 or "
"age = 47 or age = 48 or age = 49 or age = 50 or age = 51 or age = 52 or "
"age = 53 or age = 54 or age = 55 or age = 56 or age = 57 or age = 58 or "
"age = 59 or age = 60 or age = 61 or age = 62 or age = 63 or age = 64 or "
"age = 65 or age = 66 or age = 67 or age = 68 or age = 69 or age = 70 or "
"age = 71 or age = 72 or age = 73 or age = 74 or age = 75 or age = 76 or "
"age = 77 or age = 78 or age = 79 or age = 80 or age = 81 or age = 82 or "
"age = 83 or age = 84 or age = 85 or age = 86 or age = 87 or age = 88 or "
"age = 89 or age = 90 or age = 91 or age = 92 or age = 93 or age = 94 or "
"age = 95 or age = 96 or age = 97 or age = 98 or age = 99 or age = 100");
ASSERT_NE(info, nullptr);
EXPECT_EQ(
info->filter_cond()->text(),
"age in (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, "
"19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, "
"37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, "
"55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, "
"73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, "
"91, 92, 93, 94, 95, 96, 97, 98, 99, 100)(FORWARD)");
}
TEST_F(EqOrRewriteTest, SimpleManyEqOrParas) {
auto info = parse(
"age = 1 or age = 2 or age = 3 or age = 4 "
"or age = 5 or age = 6 or (age = 7 or age = 8 or age = 9 or age = 10 or "
"age = 11 or age = 12 or age = 13) or age = 14 or age = 15 or age = 16 "
"or "
"age = 17 or age = 18 or age = 19 or age = 20 or age = 21 or age = 22 or "
"age = 23 or age = 24 or age = 25 or age = 26 or age = 27 or age = 28 or "
"age = 29 or age = 30 or age = 31 or age = 32 or age = 33 or age = 34 or "
"age = 35 or age = 36 or age = 37 or (age = 38 or age = 39 or age = 40 "
"or "
"age = 41 or age = 42 or age = 43 or age = 44 or age = 45 or age = 46 or "
"age = 47 or age = 48 or age = 49 or age = 50 or age = 51 or age = 52 or "
"age = 53 or age = 54 or age = 55 or age = 56 or age = 57 or age = 58 or "
"age = 59 or age = 60 or age = 61 or age = 62 or age = 63 or age = 64 or "
"age = 65 or age = 66 or age = 67 or age = 68 or age = 69 or age = 70 or "
"age = 71 or age = 72 or age = 73 or age = 74 or age = 75 or age = 76 or "
"age = 77 or age = 78 or age = 79 or age = 80 or age = 81 or age = 82 or "
"age = 83 or age = 84 or age = 85) or age = 86 or age = 87 or age = 88 "
"or "
"age = 89 or age = 90 or age = 91 or age = 92 or age = 93 or age = 94 or "
"age = 95 or age = 96 or age = 97 or (age = 98 or age = 99) or age = "
"100");
ASSERT_NE(info, nullptr);
EXPECT_EQ(
info->filter_cond()->text(),
"age in (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, "
"19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, "
"37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, "
"55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, "
"73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, "
"91, 92, 93, 94, 95, 96, 97, 98, 99, 100)(FORWARD)");
}
TEST_F(EqOrRewriteTest, SimpleNeOr) {
auto info = parse("age != 10 or age != 20 ");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(), "age in NOT (10, 20)(FORWARD)");
}
TEST_F(EqOrRewriteTest, SimpleManyNeOr) {
auto info = parse(
"age != 1 or age != 2 or age != 3 or age "
"!= 4 or age != 5 or age != 6 or age != 7 or age != 8 or age != 9 or age "
"!= 10 or age != 11 or age != 12 or age != 13 or age != 14 or age != 15 "
"or age != 16 or age != 17 or age != 18 or age != 19 or age != 20");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"age in NOT (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, "
"16, 17, 18, "
"19, 20)(FORWARD)");
}
TEST_F(EqOrRewriteTest, EqAndNe) {
auto info = parse(
"age != 10 or age != 20 or age = 30 or "
"age = 40");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"(age in NOT (10, 20)(FORWARD)(OR_A)) or (age in (30, "
"40)(FORWARD)(OR_A))");
}
TEST_F(EqOrRewriteTest, PreEqOr) {
{
auto info = parse(
"gender =1 or age = 10 or age = 20 or "
"age = 30 or age = 40");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"(gender=1(FORWARD)(OR_A)) or (age in (10, 20, 30, "
"40)(FORWARD)(OR_A))");
}
{
auto info = parse(
"gender =1 and age = 10 or age = 20 or "
"age = 30 or age = 40");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"((gender=1(FORWARD)(OR_A)) and (age=10(FORWARD)(OR_A))) or (age "
"in (20, 30, 40)(FORWARD)(OR_A))");
}
}
TEST_F(EqOrRewriteTest, PostEqOr) {
{
auto info = parse(
"age = 10 or age = 20 or "
"age = 30 or age = 40 or gender = 1");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"(age in (10, 20, 30, 40)(FORWARD)(OR_A)) or "
"(gender=1(FORWARD)(OR_A))");
}
{
auto info = parse(
"age = 10 or age = 20 or "
"age = 30 or age = 40 and gender = 1");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"(age in (10, 20, 30)(FORWARD)(OR_A)) or "
"((age=40(FORWARD)(OR_A)) and (gender=1(FORWARD)(OR_A)))");
}
}
TEST_F(EqOrRewriteTest, PreEqAnd) {
auto info = parse(
"gender =1 and (age = 10 or age = 20 or "
"age = 30 or age = 40)");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"(gender=1(FORWARD)) and (age in (10, 20, 30, 40)(FORWARD))");
}
TEST_F(EqOrRewriteTest, PostEqAnd) {
auto info = parse(
"(age = 10 or age = 20 or "
"age = 30 or age = 40) and gender=1");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"(age in (10, 20, 30, 40)(FORWARD)) and (gender=1(FORWARD))");
}
TEST_F(EqOrRewriteTest, PrePostEqAnd) {
auto info = parse(
"gender =1 and (age = 10 or age = 20 or "
"age = 30 or age = 40) and loc != 3");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"((gender=1(FORWARD)) and (age in (10, 20, 30, 40)(FORWARD))) and "
"(loc!=3(FORWARD))");
}
TEST_F(EqOrRewriteTest, UserCases1) {
auto info = parse(
"(agent_id=20) and state=1 and (fid=107 "
"or fid=174 or fid=593 or fid=602 or fid=592 or fid=134 or fid=135 or "
"fid=136 or fid=137 or fid=138 or fid=139 or fid=141 or fid=267 or "
"fid=271 or fid=176 or fid=177 or fid=178 or fid=179 or fid=180 or "
"fid=182 or fid=183 or fid=184 or fid=270 or fid=479 or fid=488 or "
"fid=502 or fid=508 or fid=522 or fid=553 or fid=554 or fid=557 or "
"fid=561 or fid=567 or fid=570 or fid=588 or fid=594 or fid=595 or "
"fid=596 or fid=597 or fid=598 or fid=603 or fid=604 or fid=605 or "
"fid=606 or fid=426 or fid=427 or fid=428 or fid=429 or fid=430 or "
"fid=431 or fid=432 or fid=433 or fid=434 or fid=435 or fid=436 or "
"fid=437 or fid=438 or fid=439 or fid=440 or fid=441 or fid=442 or "
"fid=443 or fid=444 or fid=445 or fid=446 or fid=447 or fid=448 or "
"fid=215 or fid=216 or fid=217 or fid=469 or fid=473 or fid=475 or "
"fid=476 or fid=477 or fid=478 or fid=524 or fid=528 or fid=529 or "
"fid=532 or fid=533 or fid=534 or fid=542 or fid=543 or fid=560 or "
"fid=243 or fid=244 or fid=245 or fid=246 or fid=247 or fid=496 or "
"fid=497 or fid=506 or fid=248 or fid=249 or fid=250 or fid=251 or "
"fid=252 or fid=494 or fid=495 or fid=507 or fid=535 or fid=536 or "
"fid=586 or fid=589 or fid=259 or fid=260 or fid=261 or fid=262 or "
"fid=263 or fid=264 or fid=265 or fid=491 or fid=492 or fid=493 or "
"fid=530 or fid=531 or fid=227 or fid=228 or fid=229 or fid=230 or "
"fid=231 or fid=232 or fid=233 or fid=235 or fid=472 or fid=487 or "
"fid=537 or fid=559 or fid=236 or fid=237 or fid=238 or fid=239 or "
"fid=240 or fid=241 or fid=242 or fid=273 or fid=546 or fid=587 or "
"fid=454 or fid=455 or fid=456 or fid=457 or fid=458 or fid=459 or "
"fid=460 or fid=461 or fid=449 or fid=450 or fid=451 or fid=452 or "
"fid=453 or fid=480 or fid=481 or fid=482 or fid=483 or fid=484 or "
"fid=489 or fid=490 or fid=538 or fid=539 or fid=540 or fid=545 or "
"fid=503 or fid=504 or fid=547 or fid=548 or fid=549 or fid=550 or "
"fid=509 or fid=510 or fid=511 or fid=512 or fid=513 or fid=523 or "
"fid=558 or fid=555 or fid=556 or fid=600 or fid=601 or fid=562 or "
"fid=563 or fid=564 or fid=565 or fid=566 or fid=591 or fid=568 or "
"fid=569 or fid=590 or fid=571 or fid=572 or fid=573 or fid=574 or "
"fid=575 or fid=701 or fid=711 or fid=713 or fid=616 or fid=617 or "
"fid=618 or fid=619 or fid=620 or fid=621 or fid=622 or fid=623 or "
"fid=624 or fid=625 or fid=626 or fid=629 or fid=672 or fid=607 or "
"fid=700 or fid=635 or fid=612 or fid=613 or fid=614 or fid=615 or "
"fid=679 or fid=670 or fid=680 or fid=681 or fid=702 or fid=706 or "
"fid=714 or fid=675 or fid=676 or fid=640 or fid=643 or fid=649 or "
"fid=653 or fid=655 or fid=657 or fid=662 or fid=703 or fid=704 or "
"fid=705 or fid=707 or fid=641 or fid=642 or fid=644 or fid=645 or "
"fid=646 or fid=647 or fid=648 or fid=709 or fid=650 or fid=651 or "
"fid=652 or fid=710 or fid=654 or fid=656 or fid=658 or fid=659 or "
"fid=660 or fid=661 or fid=663 or fid=664 or fid=665 or fid=666 or "
"fid=667 or fid=668 or fid=669 or fid=678)");
ASSERT_NE(info, nullptr);
EXPECT_EQ(
info->filter_cond()->text(),
"((agent_id=20(FORWARD)) and (state=1(FORWARD))) and (fid in (107, 174, "
"593, 602, 592, 134, 135, 136, 137, 138, 139, 141, 267, 271, 176, 177, "
"178, 179, 180, 182, 183, 184, 270, 479, 488, 502, 508, 522, 553, 554, "
"557, 561, 567, 570, 588, 594, 595, 596, 597, 598, 603, 604, 605, 606, "
"426, 427, 428, 429, 430, 431, 432, 433, 434, 435, 436, 437, 438, 439, "
"440, 441, 442, 443, 444, 445, 446, 447, 448, 215, 216, 217, 469, 473, "
"475, 476, 477, 478, 524, 528, 529, 532, 533, 534, 542, 543, 560, 243, "
"244, 245, 246, 247, 496, 497, 506, 248, 249, 250, 251, 252, 494, 495, "
"507, 535, 536, 586, 589, 259, 260, 261, 262, 263, 264, 265, 491, 492, "
"493, 530, 531, 227, 228, 229, 230, 231, 232, 233, 235, 472, 487, 537, "
"559, 236, 237, 238, 239, 240, 241, 242, 273, 546, 587, 454, 455, 456, "
"457, 458, 459, 460, 461, 449, 450, 451, 452, 453, 480, 481, 482, 483, "
"484, 489, 490, 538, 539, 540, 545, 503, 504, 547, 548, 549, 550, 509, "
"510, 511, 512, 513, 523, 558, 555, 556, 600, 601, 562, 563, 564, 565, "
"566, 591, 568, 569, 590, 571, 572, 573, 574, 575, 701, 711, 713, 616, "
"617, 618, 619, 620, 621, 622, 623, 624, 625, 626, 629, 672, 607, 700, "
"635, 612, 613, 614, 615, 679, 670, 680, 681, 702, 706, 714, 675, 676, "
"640, 643, 649, 653, 655, 657, 662, 703, 704, 705, 707, 641, 642, 644, "
"645, 646, 647, 648, 709, 650, 651, 652, 710, 654, 656, 658, 659, 660, "
"661, 663, 664, 665, 666, 667, 668, 669, 678)(FORWARD))");
}
TEST_F(EqOrRewriteTest, UserCases2) {
auto info = parse(
"partitionName = '114634' or "
"partitionName = '114632' or partitionName = '114635' or partitionName = "
"'114629' or partitionName = '114630' or partitionName = '114633' or "
"partitionName = '114636' or partitionName = '114637' or partitionName = "
"'114631'");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"partitionName in (114634, 114632, 114635, 114629, 114630, 114633, "
"114636, 114637, 114631)(FORWARD)");
}
TEST_F(EqOrRewriteTest, UserCases3) {
auto info = parse(
"(doc_id=1319620650600837120 or "
"doc_id=1319621497753739264 or doc_id=1319629144649367552 or "
"doc_id=1319630319721377793 or doc_id=1319667286769324032 or "
"doc_id=1319671157117808640 or doc_id=1319671403998793728 or "
"doc_id=2319684930499055617 or doc_id=1319685259995140096)");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"doc_id in (1319620650600837120, 1319621497753739264, "
"1319629144649367552, 1319630319721377793, 1319667286769324032, "
"1319671157117808640, 1319671403998793728, 2319684930499055617, "
"1319685259995140096)(FORWARD)");
}
TEST_F(EqOrRewriteTest, UserCases4) {
auto info = parse(
"(strAttr ='' or strAttr = 'prd') and "
"categoryId = 4");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"(strAttr in (, prd)(FORWARD)) and (categoryId=4(FORWARD))");
}
TEST_F(EqOrRewriteTest, UserCases5) {
auto info = parse(
"intAttr = 1 OR intAttr = 5 OR intAttr "
"= 6 OR intAttr = 9 and categoryId = 1");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"(intAttr in (1, 5, 6)(FORWARD)(OR_A)) or "
"((intAttr=9(FORWARD)(OR_A)) and (categoryId=1(FORWARD)(OR_A)))");
}
TEST_F(EqOrRewriteTest, UserCases6) {
auto info = parse(
""
"filename='OhbVrpoi.pdf' or "
"filename='wRyoG4dB.pdf' or "
"filename='dJ3fawFf.pdf' or "
"filename='ZJS9dk3Q.pdf' or "
"filename='fY2JD8dL.pdf' or "
"filename='HnJpdoxC.pdf' or "
"filename='Hbxm1zvi.pdf' or "
"filename='r5Q8cxHu.pdf' or "
"filename='dwF9cZtI.pdf'");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"filename in (OhbVrpoi.pdf, "
"wRyoG4dB.pdf, "
"dJ3fawFf.pdf, "
"ZJS9dk3Q.pdf, "
"fY2JD8dL.pdf, "
"HnJpdoxC.pdf, "
"Hbxm1zvi.pdf, "
"r5Q8cxHu.pdf, "
"dwF9cZtI.pdf)(FORWARD)");
}
TEST_F(EqOrRewriteTest, NotChanged1) {
auto info = parse(
"passed_days>3 and (loc >= "
"500 or age > 10)");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"(passed_days>3(FORWARD)) and ((loc>=500(FORWARD)(OR_A)) "
"or (age>10(FORWARD)(OR_A)))");
}
TEST_F(EqOrRewriteTest, NotChanged2) {
auto info = parse(
"strAttr=\"online_252\" AND (intAttr > "
"103775813 OR intAttr < 103775813) and categoryId = 88888888");
ASSERT_NE(info, nullptr);
EXPECT_EQ(
info->filter_cond()->text(),
"((strAttr=online_252(FORWARD)) and ((intAttr>103775813(FORWARD)(OR_A)) "
"or (intAttr<103775813(FORWARD)(OR_A)))) and "
"(categoryId=88888888(FORWARD))");
}
TEST_F(EqOrRewriteTest, NotChanged3) {
auto info = parse(
"(is_type1 = true or is_type2 = "
"true)");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"(is_type1=true(FORWARD)(OR_A)) or (is_type2=true(FORWARD)(OR_A))");
}
TEST_F(EqOrRewriteTest, NotChanged4) {
auto info = parse("(a = 1 or a != 2)");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"(a=1(FORWARD)(OR_A)) or (a!=2(FORWARD)(OR_A))");
}
class ContainRewriteTest : public SimpleRewriterTest {};
TEST_F(ContainRewriteTest, ContainAllEmptySet) {
auto info = parse("category_array contain_all ()");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"category_array IS_NOT_NULL (FORWARD)");
}
TEST_F(ContainRewriteTest, NotContainAllEmptySet) {
auto info = parse("category_array not contain_all ()");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->is_filter_unsatisfiable(), true);
}
TEST_F(ContainRewriteTest, NotContainAnyEmptySet) {
auto info = parse("category_array not contain_any ()");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(),
"category_array IS_NOT_NULL (FORWARD)");
}
TEST_F(ContainRewriteTest, ContainAnyEmptySet) {
auto info = parse("category_array contain_any ()");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->is_filter_unsatisfiable(), true);
}
TEST_F(ContainRewriteTest, AlwaysFalseConditionAnd) {
auto info = parse("category_array not contain_all () and a = 1");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->is_filter_unsatisfiable(), true);
}
TEST_F(ContainRewriteTest, AlwaysFalseConditionMultiAnd) {
auto info = parse(
"category_array not contain_all () and a > 1 and a > 2 and a > 3 and a > "
"4");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->is_filter_unsatisfiable(), true);
}
TEST_F(ContainRewriteTest, AlwaysFalseConditionOr) {
auto info = parse("category_array not contain_all () or a = 1");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->filter_cond()->text(), "a=1(FORWARD)");
}
TEST_F(ContainRewriteTest, AlwaysFalseConditionMultiOr) {
auto info =
parse("category_array not contain_all () or a > 1 or a > 2 or a > 3");
ASSERT_NE(info, nullptr);
EXPECT_EQ(
info->filter_cond()->text(),
"((a>1(FORWARD)(OR_A)) or (a>2(FORWARD)(OR_A))) or (a>3(FORWARD)(OR_A))");
}
TEST_F(ContainRewriteTest, AlwaysFalseConditionAndComplex) {
auto info = parse("(a > 1 or a < 0) and category_array contain_any () ");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->is_filter_unsatisfiable(), true);
}
TEST_F(ContainRewriteTest, AlwaysFalseConditionOrComplex) {
auto info = parse("(a > 1 or a < 0) or category_array contain_any () ");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->is_filter_unsatisfiable(), false);
EXPECT_EQ(info->filter_cond()->text(),
"(a>1(FORWARD)(OR_A)) or (a<0(FORWARD)(OR_A))");
}
TEST_F(SimpleRewriterTest, MiscOr) {
auto info = parse("a = 1 or a = 2 or a = 3 or category_array contain_any ()");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->is_filter_unsatisfiable(), false);
EXPECT_EQ(info->filter_cond()->text(), "a in (1, 2, 3)(FORWARD)");
}
TEST_F(SimpleRewriterTest, MiscAnd) {
auto info =
parse("(a = 1 or a = 2 or a = 3) and category_array contain_any ()");
ASSERT_NE(info, nullptr);
EXPECT_EQ(info->is_filter_unsatisfiable(), true);
}
} // namespace zvec::sqlengine
+171
View File
@@ -0,0 +1,171 @@
// 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 "db/sqlengine/sqlengine.h"
#include <cstdint>
#include <memory>
#include <gtest/gtest.h>
#include "zvec/db//schema.h"
#include "zvec/db/query_params.h"
#include "zvec/db/type.h"
#include "mock_segment.h"
namespace zvec::sqlengine {
class SqlEngineTest : public testing::Test {
public:
void SetUp() override {
auto vector_params = std::make_shared<FlatIndexParams>(MetricType::IP);
schema_ = std::make_shared<CollectionSchema>(
"test_collection",
std::vector<FieldSchema::Ptr>{
std::make_shared<FieldSchema>("id", DataType::INT32, false, 0,
nullptr),
std::make_shared<FieldSchema>(
"name", DataType::STRING, false, 0, // nullptr
std::make_shared<InvertIndexParams>(false)),
std::make_shared<FieldSchema>("age", DataType::INT64, false, 0,
nullptr),
std::make_shared<FieldSchema>("score", DataType::DOUBLE, false, 0,
nullptr),
std::make_shared<FieldSchema>("tag_list", DataType::ARRAY_INT32,
false, 0, nullptr),
std::make_shared<FieldSchema>("vector",
DataType::SPARSE_VECTOR_FP32, false,
4, vector_params),
});
}
protected:
CollectionSchema::Ptr schema_;
};
TEST_F(SqlEngineTest, Forward) {
std::vector<Segment::Ptr> segments = {std::make_shared<MockSegment>()};
SearchQuery query;
query.output_fields_ = {"id", "name", "age", "tag_list"};
query.topk_ = 11;
// query.filter_ = "id > 3 and score < 0.1";
// query.filter_ = "name like 'name_2%'";
// query.filter_ = "name not in ('name_2','name_4')";
// query.filter_ = "tag_list contain_all (1,2,3,4)";
query.filter_ = "tag_list is null";
if (const char *env_var = std::getenv("FILTER"); env_var != nullptr) {
query.filter_ = env_var;
}
query.include_vector_ = true;
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(schema_, query, segments);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
EXPECT_TRUE(ret.has_value());
}
TEST_F(SqlEngineTest, Vector) {
std::vector<Segment::Ptr> segments = {std::make_shared<MockSegment>()};
SearchQuery query;
query.output_fields_ = {"id", "name", "score"};
query.topk_ = 11;
query.filter_ = "id > 3 and score < 0.1";
if (const char *env_var = std::getenv("FILTER"); env_var != nullptr) {
query.filter_ = env_var;
}
// query.target_.set_vector("[0.1, 0.2, 0.3, 0.4]");
query.target_.set_sparse_vector("[0, 1, 2, 3]", "[0.1, 0.2, 0.3, 0.4]");
query.target_.field_name_ = "vector";
query.include_vector_ = true;
query.target_.query_params_ = std::make_shared<QueryParams>(IndexType::FLAT);
query.target_.query_params_->set_radius(0.8F);
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(schema_, query, segments);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
EXPECT_TRUE(ret.has_value());
}
TEST_F(SqlEngineTest, Invert) {
std::vector<Segment::Ptr> segments = {std::make_shared<MockSegment>()};
SearchQuery query;
query.output_fields_ = {"id", "age", "score"};
query.topk_ = 11;
// query.filter_ = "name = 'test_name'";
query.filter_ = "name is not null";
if (const char *env_var = std::getenv("FILTER"); env_var != nullptr) {
query.filter_ = env_var;
}
query.include_vector_ = true;
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(schema_, query, segments);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
EXPECT_TRUE(ret.has_value());
}
TEST_F(SqlEngineTest, MultiSegments) {
std::vector<Segment::Ptr> segments = {std::make_shared<MockSegment>(),
std::make_shared<MockSegment>()};
SearchQuery query;
query.output_fields_ = {"id", "name", "age", "score"};
query.topk_ = 11;
query.target_.set_vector("[0.1, 0.2, 0.3, 0.4]");
query.target_.field_name_ = "vector";
// query.filter_ = "name = 'test_name'";
if (const char *env_var = std::getenv("FILTER"); env_var != nullptr) {
query.filter_ = env_var;
}
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(schema_, query, segments);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
EXPECT_TRUE(ret.has_value());
}
TEST_F(SqlEngineTest, GroupBy) {
std::vector<Segment::Ptr> segments = {std::make_shared<MockSegment>()};
GroupByVectorQuery query;
query.group_by_field_name_ = "name";
query.group_count_ = 3;
query.group_topk_ = 2;
query.output_fields_ = {"id", "name", "score"};
query.filter_ = "id > 3 and score < 0.1";
if (const char *env_var = std::getenv("FILTER"); env_var != nullptr) {
query.filter_ = env_var;
}
// query.target_.set_vector("[0.1, 0.2, 0.3, 0.4]");
query.target_.set_sparse_vector("[0, 1, 2, 3]", "[0.1, 0.2, 0.3, 0.4]");
query.target_.field_name_ = "vector";
query.include_vector_ = true;
query.target_.query_params_ = std::make_shared<QueryParams>(IndexType::FLAT);
query.target_.query_params_->set_radius(0.8F);
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute_group_by(schema_, query, segments);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
EXPECT_TRUE(ret.has_value());
}
} // namespace zvec::sqlengine
+112
View File
@@ -0,0 +1,112 @@
// 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
#pragma once
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <memory>
#include <arrow/api.h>
#include <arrow/io/api.h>
#include <arrow/ipc/api.h>
#include <gtest/gtest.h>
#include "db/common/file_helper.h"
#include "db/index/common/version_manager.h"
#include "db/index/segment/segment.h"
#include "db/sqlengine/sqlengine.h"
#include "zvec/db/index_params.h"
#include "zvec/db/schema.h"
#include "zvec/db/type.h"
namespace zvec::sqlengine {
using CreateDocFun = Doc (*)(const uint64_t doc_id);
inline Status InsertDoc(const Segment::Ptr &segment,
const uint64_t start_doc_id, const uint64_t end_doc_id,
CreateDocFun create_doc) {
srand(time(nullptr));
long long create_total = 0;
long long insert_total = 0;
for (auto doc_id = start_doc_id; doc_id < end_doc_id; doc_id++) {
if (segment) {
auto start = std::chrono::system_clock::now();
Doc new_doc = create_doc(doc_id);
auto end = std::chrono::system_clock::now();
auto create_cost =
std::chrono::duration_cast<std::chrono::microseconds>(end - start)
.count();
create_total += create_cost;
start = std::chrono::system_clock::now();
auto status = segment->Insert(new_doc);
if (!status.ok()) {
return status;
}
end = std::chrono::system_clock::now();
auto insert_cost =
std::chrono::duration_cast<std::chrono::microseconds>(end - start)
.count();
insert_total += insert_cost;
}
}
std::cout << "pure create cost " << create_total << "us" << std::endl;
std::cout << "pure insert cost " << insert_total << "us" << std::endl;
return Status::OK();
}
inline Segment::Ptr create_segment(const std::string &seg_path,
const CollectionSchema &schema) {
auto segment_meta = std::make_shared<SegmentMeta>();
segment_meta->set_id(0);
auto id_map = IDMap::CreateAndOpen("test_collection", seg_path + "/id_map",
true, false);
auto delete_store = std::make_shared<DeleteStore>("test_collection");
Version v1;
v1.set_schema(schema);
std::string v_path = seg_path + "/test_manifest";
FileHelper::CreateDirectory(v_path);
auto vm = VersionManager::Create(v_path, v1);
if (!vm.has_value()) {
LOG_ERROR("create version manager failed: %s", vm.error().c_str());
return nullptr;
}
BlockMeta mem_block;
mem_block.id_ = 0;
mem_block.type_ = BlockType::SCALAR;
mem_block.min_doc_id_ = 0;
mem_block.max_doc_id_ = 0;
mem_block.doc_count_ = 0;
segment_meta->set_writing_forward_block(mem_block);
SegmentOptions options;
options.read_only_ = false;
options.enable_mmap_ = true;
auto result = Segment::CreateAndOpen(seg_path, schema, 0, 0, id_map,
delete_store, vm.value(), options);
if (!result) {
LOG_ERROR("create segment failed: %s", result.error().c_str());
return nullptr;
}
auto segment = result.value();
return segment;
}
} // namespace zvec::sqlengine
+291
View File
@@ -0,0 +1,291 @@
// 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 <cstdint>
#include <memory>
#include <gtest/gtest.h>
#include "db/sqlengine/sqlengine.h"
#include "recall_base.h"
namespace zvec::sqlengine {
class VectorRecallTest : public RecallTest {};
TEST_F(VectorRecallTest, Basic) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
std::vector<float> feature(4, 0.0);
query.target_.set_vector(std::string((const char *)feature.data(),
feature.size() * sizeof(float)));
query.target_.field_name_ = "dense";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value());
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int i = 0; i < query.topk_; i++) {
auto &doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
EXPECT_FLOAT_EQ(doc->score(), (float)i * i * 4);
}
}
TEST_F(VectorRecallTest, HybridInvertFilter) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.filter_ = "invert_id >= 1";
query.topk_ = 200;
std::vector<float> feature(4, 0.0);
query.target_.set_vector(std::string((const char *)feature.data(),
feature.size() * sizeof(float)));
query.target_.field_name_ = "dense";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value());
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int j = 0; j < query.topk_; j++) {
auto &doc = docs[j];
int i = j + 1;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
EXPECT_FLOAT_EQ(doc->score(), (float)i * i * 4);
}
}
TEST_F(VectorRecallTest, HybridInvertFilterBfByKeys) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.filter_ = "invert_id < 199";
query.topk_ = 199;
std::vector<float> feature(4, 0.0);
query.target_.set_vector(std::string((const char *)feature.data(),
feature.size() * sizeof(float)));
query.target_.field_name_ = "dense";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value());
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int i = 0; i < query.topk_; i++) {
auto &doc = docs[i];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
EXPECT_FLOAT_EQ(doc->score(), (float)i * i * 4);
}
}
TEST_F(VectorRecallTest, HybridForwardFilter) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.filter_ = "id >= 1";
query.topk_ = 200;
std::vector<float> feature(4, 0.0);
query.target_.set_vector(std::string((const char *)feature.data(),
feature.size() * sizeof(float)));
query.target_.field_name_ = "dense";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value());
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
for (int j = 0; j < query.topk_; j++) {
auto &doc = docs[j];
int i = j + 1;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(i));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), i % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(i % 100));
EXPECT_FLOAT_EQ(doc->score(), (float)i * i * 4);
}
}
TEST_F(VectorRecallTest, HybridInvertForwardFilter) {
SearchQuery query;
query.output_fields_ = {"name", "age"};
query.filter_ = "invert_id >= 1 and id <= 100";
query.topk_ = 200;
std::vector<float> feature(4, 0.0);
query.target_.set_vector(std::string((const char *)feature.data(),
feature.size() * sizeof(float)));
query.target_.field_name_ = "dense";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value());
auto docs = ret.value();
EXPECT_EQ(docs.size(), 100);
for (size_t j = 0; j < docs.size(); j++) {
auto &doc = docs[j];
int doc_id = j + 1;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), doc_id % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(doc_id % 100));
EXPECT_FLOAT_EQ(doc->score(), (float)doc_id * doc_id * 4);
}
}
TEST_F(VectorRecallTest, Sparse) {
SearchQuery query;
query.output_fields_ = {"id", "name", "age"};
query.topk_ = 200;
std::vector<float> feature(4, 1.0);
std::vector<uint32_t> indices{0, 1, 2, 3};
query.target_.set_sparse_vector(
std::string((const char *)indices.data(),
indices.size() * sizeof(uint32_t)),
std::string((const char *)feature.data(),
feature.size() * sizeof(float)));
query.target_.field_name_ = "sparse";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value());
auto docs = ret.value();
EXPECT_EQ(docs.size(), query.topk_);
int doc_id = 9999;
for (size_t j = 0; j < docs.size(); j++) {
auto &doc = docs[j];
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), doc_id % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(doc_id % 100));
EXPECT_FLOAT_EQ(doc->score(), (float)doc_id * 4);
doc_id--;
while (doc_id % 100 <= 3) {
doc_id--;
}
}
}
TEST_F(VectorRecallTest, DeleteFilter) {
// This test uses only one segment and thus we only operate on the first one
for (int i = 0; i < 4000; i++) {
segments_[0]->Delete("pk_" + std::to_string(i));
}
SearchQuery query;
query.output_fields_ = {"name", "age"};
query.topk_ = 100;
std::vector<float> feature(4, 0.0);
query.target_.set_vector(std::string((const char *)feature.data(),
feature.size() * sizeof(float)));
query.target_.field_name_ = "dense";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value());
auto docs = ret.value();
EXPECT_EQ(docs.size(), 100);
for (size_t j = 0; j < docs.size(); j++) {
auto &doc = docs[j];
int doc_id = j + 4000;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), doc_id % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(doc_id % 100));
EXPECT_FLOAT_EQ(doc->score(), (float)doc_id * doc_id * 4);
}
}
TEST_F(VectorRecallTest, HybridInvertForwardDeleteFilter) {
// In previous test, docs[0-4000) has been deleted
SearchQuery query;
query.output_fields_ = {"name", "age"};
query.filter_ = "invert_id >= 6000 and id < 6080";
query.topk_ = 100;
std::vector<float> feature(4, 0.0);
query.target_.set_vector(std::string((const char *)feature.data(),
feature.size() * sizeof(float)));
query.target_.field_name_ = "dense";
auto engine = SQLEngine::create(std::make_shared<Profiler>());
auto ret = engine->execute(collection_schema_, query, segments_);
if (!ret) {
LOG_ERROR("execute failed: [%s]", ret.error().c_str());
}
ASSERT_TRUE(ret.has_value());
auto docs = ret.value();
EXPECT_EQ(docs.size(), 80);
for (size_t j = 0; j < docs.size(); j++) {
auto &doc = docs[j];
int doc_id = j + 6000;
EXPECT_EQ(doc->pk(), "pk_" + std::to_string(doc_id));
auto age = doc->get<int32_t>("age");
EXPECT_EQ(age.value(), doc_id % 100);
auto name = doc->get<std::string>("name");
ASSERT_TRUE(name);
EXPECT_EQ(name.value(), "user_" + std::to_string(doc_id % 100));
EXPECT_FLOAT_EQ(doc->score(), (float)doc_id * doc_id * 4);
}
}
} // namespace zvec::sqlengine
+211
View File
@@ -0,0 +1,211 @@
// 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 <cstdio>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include <zvec/ailego/io/file.h>
#include <zvec/ailego/utility/file_helper.h>
#include "db/common/file_helper.h"
#include "index/utils/utils.h"
#include "zvec/db/collection.h"
#include "zvec/db/doc.h"
#include "zvec/db/options.h"
#include "zvec/db/schema.h"
#include "zvec/db/status.h"
#if defined(_WIN32) || defined(_WIN64)
#define WIN32_LEAN_AND_MEAN
#define NOMINMAX
#include <windows.h>
#undef DeleteFile
#undef RemoveDirectory
#undef CreateFile
#undef DELETE
#endif
using namespace zvec;
using namespace zvec::test;
// UTF-8 test path containing Chinese characters.
// If the path handling is broken, the OS will show garbled names (mojibake).
static const std::string kUtf8Dir =
"utf8_col_\xe4\xb8\xad\xe6\x96\x87\xe6\xb5\x8b\xe8\xaf\x95"; // utf8_col_中文测试
class Utf8CollectionTest : public ::testing::Test {
protected:
void SetUp() override {
ailego::FileHelper::RemovePath(kUtf8Dir.c_str());
}
void TearDown() override {
ailego::FileHelper::RemovePath(kUtf8Dir.c_str());
}
};
// ---------------------------------------------------------------------------
// 1. Create a collection in a UTF-8 path, insert docs, flush, reopen
// ---------------------------------------------------------------------------
TEST_F(Utf8CollectionTest, CreateInsertFlushReopen) {
CollectionOptions opts;
opts.read_only_ = false;
opts.enable_mmap_ = true;
auto schema = TestHelper::CreateNormalSchema();
auto result = Collection::CreateAndOpen(kUtf8Dir, *schema, opts);
if (!result.has_value()) {
std::cout << result.error().message() << std::endl;
}
ASSERT_TRUE(result.has_value());
ASSERT_TRUE(ailego::FileHelper::IsExist(kUtf8Dir.c_str()));
auto col = std::move(result).value();
ASSERT_EQ(col->Path(), kUtf8Dir);
const int kDocCount = 10;
auto s = TestHelper::CollectionInsertDoc(col, 0, kDocCount);
ASSERT_TRUE(s.ok()) << s.message();
ASSERT_TRUE(col->Flush().ok());
auto stats = col->Stats().value();
ASSERT_EQ(stats.doc_count, kDocCount);
col.reset();
// Reopen and verify doc count survives
auto reopen = Collection::Open(kUtf8Dir, opts);
ASSERT_TRUE(reopen.has_value()) << reopen.error().message();
auto col2 = std::move(reopen).value();
auto stats2 = col2->Stats().value();
ASSERT_EQ(stats2.doc_count, kDocCount);
}
// ---------------------------------------------------------------------------
// 2. Destroy a collection in a UTF-8 path
// ---------------------------------------------------------------------------
TEST_F(Utf8CollectionTest, CreateAndDestroy) {
CollectionOptions opts;
opts.read_only_ = false;
opts.enable_mmap_ = true;
auto schema = TestHelper::CreateNormalSchema();
auto result = Collection::CreateAndOpen(kUtf8Dir, *schema, opts);
ASSERT_TRUE(result.has_value());
auto col = std::move(result).value();
ASSERT_EQ(col->Destroy(), Status::OK());
ASSERT_FALSE(ailego::FileHelper::IsExist(kUtf8Dir.c_str()));
}
// ---------------------------------------------------------------------------
// 3. Verify the on-disk layout contains correct UTF-8 names (no mojibake)
// by listing files via the wide-char Windows API or POSIX readdir.
// ---------------------------------------------------------------------------
#if defined(_WIN32) || defined(_WIN64)
TEST_F(Utf8CollectionTest, NoDirGarble) {
CollectionOptions opts;
opts.read_only_ = false;
opts.enable_mmap_ = true;
auto schema = TestHelper::CreateNormalSchema();
auto result = Collection::CreateAndOpen(kUtf8Dir, *schema, opts);
ASSERT_TRUE(result.has_value());
auto col = std::move(result).value();
const int kDocCount = 5;
auto s = TestHelper::CollectionInsertDoc(col, 0, kDocCount);
ASSERT_TRUE(s.ok()) << s.message();
ASSERT_TRUE(col->Flush().ok());
// --- Check parent directory via FindFirstFileW ---
{
WIN32_FIND_DATAW fd;
HANDLE hFind = FindFirstFileW(L".\\*", &fd);
ASSERT_NE(INVALID_HANDLE_VALUE, hFind);
std::wstring expected = ailego::FileHelper::Utf8ToWide(kUtf8Dir);
bool found = false;
do {
if (std::wstring(fd.cFileName) == expected) {
found = true;
break;
}
} while (FindNextFileW(hFind, &fd));
FindClose(hFind);
EXPECT_TRUE(found) << "Collection dir '" << kUtf8Dir
<< "' not found via FindFirstFileW (garbled Unicode?)";
}
// --- Check inside collection via "cmd /u /c dir /b" redirected to file ---
{
std::wstring wdir = ailego::FileHelper::Utf8ToWide(kUtf8Dir);
std::string tmpfile = "utf8_col_dir_output.txt";
std::wstring wtmp = ailego::FileHelper::Utf8ToWide(tmpfile);
std::wstring cmd = L"cmd /u /c dir /b \"" + wdir + L"\" >\"" + wtmp + L"\"";
_wsystem(cmd.c_str());
std::ifstream ifs;
ailego::FileHelper::OpenIfstream(ifs, tmpfile, std::ios::binary);
ASSERT_TRUE(ifs.is_open());
std::string raw((std::istreambuf_iterator<char>(ifs)),
std::istreambuf_iterator<char>());
ifs.close();
ailego::FileHelper::DeleteFile(tmpfile.c_str());
const wchar_t *wdata = reinterpret_cast<const wchar_t *>(raw.data());
size_t wlen = raw.size() / sizeof(wchar_t);
if (wlen > 0 && wdata[0] == L'\xFEFF') {
wdata++;
wlen--;
}
std::wstring woutput(wdata, wlen);
std::string output = ailego::FileHelper::WideToUtf8(woutput);
EXPECT_FALSE(output.empty()) << "dir listing inside collection is empty";
std::cout << "[NoDirGarble] Contents of " << kUtf8Dir << ":\n"
<< output << std::endl;
}
}
#endif
// ---------------------------------------------------------------------------
// 4. MakeWalPath / MakeSegmentPath etc. produce correct UTF-8 results
// ---------------------------------------------------------------------------
TEST_F(Utf8CollectionTest, FileHelperPaths) {
std::string wal = FileHelper::MakeWalPath(kUtf8Dir, 0, 1);
std::string seg = FileHelper::MakeSegmentPath(kUtf8Dir, 0);
std::string fwd = FileHelper::MakeForwardBlockPath(kUtf8Dir, 0, 1, true);
std::string inv = FileHelper::MakeInvertIndexPath(kUtf8Dir, 0, 1);
std::string vec = FileHelper::MakeVectorIndexPath(kUtf8Dir, "vec1", 0, 1);
// All paths should start with our UTF-8 base
EXPECT_EQ(0u, wal.find(kUtf8Dir));
EXPECT_EQ(0u, seg.find(kUtf8Dir));
EXPECT_EQ(0u, fwd.find(kUtf8Dir));
EXPECT_EQ(0u, inv.find(kUtf8Dir));
EXPECT_EQ(0u, vec.find(kUtf8Dir));
// Should not contain forward slash on Windows
#if defined(_WIN32) || defined(_WIN64)
EXPECT_EQ(std::string::npos, wal.find('/'));
EXPECT_EQ(std::string::npos, seg.find('/'));
EXPECT_EQ(std::string::npos, fwd.find('/'));
EXPECT_EQ(std::string::npos, inv.find('/'));
EXPECT_EQ(std::string::npos, vec.find('/'));
#endif
}