wehub/alibaba--zvec
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
tests/db/sqlengine/CMakeLists.txt
+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()
tests/db/sqlengine/contain_test.cc
+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
tests/db/sqlengine/forward_recall_test.cc
+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
tests/db/sqlengine/fts_multi_segment_test.cc
+233
View File
@@ -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
tests/db/sqlengine/fts_parser_test.cc
+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
tests/db/sqlengine/fts_recall_test.cc
+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
tests/db/sqlengine/invert_recall_test.cc
+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
tests/db/sqlengine/like_test.cc
+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
tests/db/sqlengine/mock_segment.h
+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
tests/db/sqlengine/optimizer_test.cc
+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
tests/db/sqlengine/query_info_test.cc
+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
tests/db/sqlengine/recall_base.h
+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
tests/db/sqlengine/simple_rewriter_test.cc
+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
tests/db/sqlengine/sqlengine_test.cc
+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
tests/db/sqlengine/test_helper.h
+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
tests/db/sqlengine/vector_recall_test.cc
+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