fix: fix parse plan proto failed for search type (#34944)

#25848

Signed-off-by: luzhang <luzhang@zilliz.com>
Co-authored-by: luzhang <luzhang@zilliz.com>
pull/35082/head
zhagnlu 2024-07-29 21:19:49 +08:00 committed by GitHub
parent d36cfe71e5
commit a8a4779749
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 45 additions and 10 deletions

View File

@ -80,13 +80,28 @@ ParsePlaceholderGroup(const Plan* plan,
return result;
}
void
ParsePlanNodeProto(proto::plan::PlanNode& plan_node,
const void* serialized_expr_plan,
int64_t size) {
google::protobuf::io::ArrayInputStream array_stream(serialized_expr_plan,
size);
google::protobuf::io::CodedInputStream input_stream(&array_stream);
input_stream.SetRecursionLimit(std::numeric_limits<int32_t>::max());
auto res = plan_node.ParsePartialFromCodedStream(&input_stream);
if (!res) {
PanicInfo(UnexpectedError, "parse plan node proto failed");
}
}
std::unique_ptr<Plan>
CreateSearchPlanByExpr(const Schema& schema,
const void* serialized_expr_plan,
const int64_t size) {
// Note: serialized_expr_plan is of binary format
proto::plan::PlanNode plan_node;
plan_node.ParseFromArray(serialized_expr_plan, size);
ParsePlanNodeProto(plan_node, serialized_expr_plan, size);
return ProtoParser(schema).CreatePlan(plan_node);
}
@ -101,15 +116,7 @@ CreateRetrievePlanByExpr(const Schema& schema,
const void* serialized_expr_plan,
const int64_t size) {
proto::plan::PlanNode plan_node;
google::protobuf::io::ArrayInputStream array_stream(serialized_expr_plan,
size);
google::protobuf::io::CodedInputStream input_stream(&array_stream);
input_stream.SetRecursionLimit(std::numeric_limits<int32_t>::max());
auto res = plan_node.ParsePartialFromCodedStream(&input_stream);
if (!res) {
PanicInfo(UnexpectedError, "parse plan node proto failed");
}
ParsePlanNodeProto(plan_node, serialized_expr_plan, size);
return ProtoParser(schema).CreateRetrievePlan(plan_node);
}

View File

@ -26,6 +26,11 @@ struct Plan;
struct PlaceholderGroup;
struct RetrievePlan;
void
ParsePlanNodeProto(proto::plan::PlanNode& plan_node,
const void* serialized_expr_plan,
int64_t size);
// Note: serialized_expr_plan is of binary format
std::unique_ptr<Plan>
CreateSearchPlanByExpr(const Schema& schema,

View File

@ -2579,6 +2579,29 @@ class TestQueryOperation(TestcaseBase):
_, check_res = collection_w.query(multi_exprs, output_fields=[f'{default_int_field_name}'])
assert(check_res == True)
@pytest.mark.tags(CaseLabel.L0)
def test_search_multi_logical_exprs(self):
"""
target: test the scenario which search with many logical expressions
method: 1. create collection
3. search with the expr that like: int64 == 0 || int64 == 1 ........
expected: run successfully
"""
c_name = cf.gen_unique_str(prefix)
collection_w = self.init_collection_wrap(name=c_name)
df = cf.gen_default_dataframe_data()
collection_w.insert(df)
collection_w.create_index(ct.default_float_vec_field_name, index_params=ct.default_flat_index)
collection_w.load()
multi_exprs = " || ".join(f'{default_int_field_name} == {i}' for i in range(60))
collection_w.load()
vectors_s = [[random.random() for _ in range(ct.default_dim)] for _ in range(ct.default_nq)]
limit = 1000
_, check_res = collection_w.search(vectors_s[:ct.default_nq], ct.default_float_vec_field_name,
ct.default_search_params, limit, multi_exprs)
assert(check_res == True)
class TestQueryString(TestcaseBase):
"""