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

pr: #34944 Signed-off-by: luzhang <[email protected]> Co-authored-by: luzhang <[email protected]>
milvus-io · Jul 24, 2024 · a6c3056 · a6c3056
1 parent fe786ff
commit a6c3056
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 10 deletions.
diff --git a/internal/core/src/query/Plan.cpp b/internal/core/src/query/Plan.cpp
@@ -68,13 +68,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);
 }
 
@@ -83,15 +98,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) {
-        throw SegcoreError(UnexpectedError, "parse plan node proto failed");
-    }
+    ParsePlanNodeProto(plan_node, serialized_expr_plan, size);
     return ProtoParser(schema).CreateRetrievePlan(plan_node);
 }
 

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

diff --git a/tests/python_client/testcases/test_query.py b/tests/python_client/testcases/test_query.py
@@ -2370,6 +2370,29 @@ def test_query_multi_logical_exprs(self):
         _, 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):
     """