feat: enlarge topK limit (#348)

This commit is contained in:
Qinren Zhou 2026-04-17 17:51:07 +08:00 committed by GitHub
parent 3f9ebbbe96
commit f602ed30ce
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 43 additions and 32 deletions

View File

@ -1,5 +1,5 @@
<p align="right">
English | <a href="./README_CN.md">简体中文</a>
English | <a href="./README_CN.md">中文</a>
</p>
<div align="center">

View File

@ -1,5 +1,5 @@
<p align="right">
<a href="./README.md">English</a> | 简体中文
<a href="./README.md">English</a> | 中文
</p>
<div align="center">

View File

@ -1168,8 +1168,8 @@ class TestDefaultLocalDenseEmbedding:
return_value="/path/to/model",
):
mock_ms = Mock()
mock_require_module.side_effect = (
lambda m: mock_st if m == "sentence_transformers" else mock_ms
mock_require_module.side_effect = lambda m: (
mock_st if m == "sentence_transformers" else mock_ms
)
emb_func_ms = DefaultLocalDenseEmbedding(model_source="modelscope")
assert (
@ -1635,8 +1635,8 @@ class TestDefaultLocalSparseEmbedding:
"modelscope.hub.snapshot_download.snapshot_download",
return_value="/cache/splade-cocondenser",
):
mock_require_module.side_effect = (
lambda m: mock_st if m == "sentence_transformers" else mock_ms
mock_require_module.side_effect = lambda m: (
mock_st if m == "sentence_transformers" else mock_ms
)
sparse_emb = DefaultLocalSparseEmbedding(model_source="modelscope")

View File

@ -52,7 +52,7 @@ constexpr uint32_t kMaxScalarFieldSize = 1024;
constexpr uint32_t kMaxVectorFieldSize = 5;
constexpr uint32_t kMaxQueryTopk = 1024;
constexpr uint32_t kMaxQueryTopk = 100000;
constexpr uint32_t kMaxOutputFieldSize = 1024;

View File

@ -1203,53 +1203,61 @@ bool Doc::operator==(const Doc &other) const {
Status VectorQuery::validate(const FieldSchema *schema) const {
if ((uint32_t)topk_ > kMaxQueryTopk) {
return Status::InvalidArgument("query validate failed: topk[", topk_,
"] is too large, max is ", kMaxQueryTopk);
return Status::InvalidArgument("Invalid query: topk[", topk_,
"] exceeds the maximum allowed value of ",
kMaxQueryTopk);
}
if (output_fields_.has_value() &&
output_fields_->size() > kMaxOutputFieldSize) {
return Status::InvalidArgument(
"query validate failed: output_fields is too large, max is ",
"Invalid query: too many output fields, the maximum allowed is ",
kMaxOutputFieldSize);
}
if (schema == nullptr) {
// support query with vector
if (query_vector_.empty() && query_sparse_indices_.empty()) {
// No vector provided — this is a scalar-only filter query.
return Status::OK();
} else {
// If a query vector was provided, the field must exist as a vector field
// since we are doing vector similarity search.
return Status::InvalidArgument(
"Invalid query: query vector is provided, but query field[",
field_name_,
"] does not exist or is not a vector field in the collection");
}
return Status::InvalidArgument("query validate failed: vector_field[",
field_name_,
"] not defined in the collection schema");
}
// validate dense/sparse vector
if (schema->is_dense_vector()) {
// validate dimension
// Validate dimension
auto dim = schema->dimension();
switch (schema->data_type()) {
case DataType::VECTOR_FP16:
if (dim * sizeof(float16_t) != query_vector_.size()) {
return Status::InvalidArgument(
"query validate failed: dimension is invalid");
"Invalid query: dimension mismatch, expected ", dim, " but got ",
query_vector_.size() / sizeof(float16_t), " (FP16)");
}
break;
case DataType::VECTOR_FP32:
if (dim * sizeof(float) != query_vector_.size()) {
return Status::InvalidArgument(
"query validate failed: dimension is invalid");
"Invalid query: dimension mismatch, expected ", dim, " but got ",
query_vector_.size() / sizeof(float), " (FP32)");
}
break;
case DataType::VECTOR_FP64:
if (dim * sizeof(double) != query_vector_.size()) {
return Status::InvalidArgument(
"query validate failed: dimension is invalid");
"Invalid query: dimension mismatch, expected ", dim, " but got ",
query_vector_.size() / sizeof(double), " (FP64)");
}
break;
case DataType::VECTOR_INT8:
if (dim * sizeof(int8_t) != query_vector_.size()) {
return Status::InvalidArgument(
"query validate failed: dimension is invalid");
"Invalid query: dimension mismatch, expected ", dim, " but got ",
query_vector_.size() / sizeof(int8_t), " (INT8)");
}
break;
case DataType::VECTOR_INT16:
@ -1257,22 +1265,22 @@ Status VectorQuery::validate(const FieldSchema *schema) const {
case DataType::VECTOR_BINARY32:
case DataType::VECTOR_BINARY64:
return Status::NotSupported(
"query validate failed: unsupported dense vector type");
"Invalid query: dense vector type of field[", field_name_,
"] is not supported");
default:
return Status::InvalidArgument(
"query validate failed: field is not dense vector");
return Status::InvalidArgument("Invalid query: field[", field_name_,
"] is not a dense vector field");
}
} else if (schema->is_sparse_vector()) {
// validate sparse indices size
// Validate sparse indices size
if (query_sparse_indices_.size() > kSparseMaxDimSize * sizeof(uint32_t)) {
return Status::InvalidArgument(
"query validate failed: the number of sparse indices exceeds the "
"maximum limit ",
"Invalid query: too many sparse indices, the maximum allowed is ",
kSparseMaxDimSize);
}
} else {
return Status::InvalidArgument(
"query validate failed: field is not vector");
return Status::InvalidArgument("Invalid query: field[", field_name_,
"] is not a vector field");
}
return Status::OK();
}

View File

@ -30,7 +30,6 @@
#include "db/index/common/type_helper.h"
#include "index/utils/utils.h"
#include "zvec/ailego/utility/float_helper.h"
#include "zvec/db/config.h"
#include "zvec/db/doc.h"
#include "zvec/db/index_params.h"
#include "zvec/db/options.h"
@ -2110,7 +2109,8 @@ TEST_F(CollectionTest, Feature_CreateIndex_Vector) {
TEST_F(CollectionTest, Feature_CreateIndex_Scalar) {
#ifdef __ANDROID__
GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink support (needed by RocksDB checkpoint)";
GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink "
"support (needed by RocksDB checkpoint)";
#endif
auto func = [&](std::string field_name, bool enable_optimize,
IndexParams::Ptr scalar_index_params = nullptr) {
@ -2387,7 +2387,8 @@ TEST_F(CollectionTest, Feature_DropIndex_Vector) {
TEST_F(CollectionTest, Feature_DropIndex_Scalar) {
#ifdef __ANDROID__
GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink support (needed by RocksDB checkpoint)";
GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink "
"support (needed by RocksDB checkpoint)";
#endif
auto func = [&](std::string field_name, bool enable_optimize) {
FileHelper::RemoveDirectory(col_path);
@ -3266,7 +3267,7 @@ TEST_F(CollectionTest, Feature_Query_Validate) {
{
VectorQuery query;
query.topk_ = 1025;
query.topk_ = 100001;
query.field_name_ = field_name;
auto field_scheama = schema->get_vector_field(field_name);

View File

@ -50,6 +50,7 @@ cc_binary(
INCS .. ../../src
LDFLAGS ${APPLE_FRAMEWORK_LIBS}
)
set_target_properties(data_generator PROPERTIES EXCLUDE_FROM_ALL TRUE)
# Build collection_optimizer executable
@ -60,6 +61,7 @@ cc_binary(
INCS .. ../../src
LDFLAGS ${APPLE_FRAMEWORK_LIBS}
)
set_target_properties(collection_optimizer PROPERTIES EXCLUDE_FROM_ALL TRUE)
# Build test executables