Merge pull request #9 from reddit/v2.6.6-patched-251123-trace

Minimum Should Match Support -cherry pick from master

Author: kumaramit01

client/milvusclient/read_example_test.go

@@ -403,3 +403,91 @@ func ExampleClient_HybridSearch() {
 		log.Println("Scores: ", resultSet.Scores)
 	}
 }
+
+func ExampleClient_Search_textMatch() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	collectionName := "text_min_match"
+	titleField := "title"
+	textField := "document_text"
+	titleSparse := "title_sparse_vector"
+	textSparse := "text_sparse_vector"
+	milvusAddr := "127.0.0.1:19530"
+
+	cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
+		Address: milvusAddr,
+	})
+	if err != nil {
+		log.Fatal("failed to connect to milvus server: ", err.Error())
+	}
+	defer cli.Close(ctx)
+
+	_ = cli.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName))
+
+	schema := entity.NewSchema().
+		WithField(entity.NewField().WithName("id").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true).WithIsAutoID(true)).
+		WithField(entity.NewField().WithName(titleField).WithDataType(entity.FieldTypeVarChar).WithMaxLength(512).WithEnableAnalyzer(true).WithEnableMatch(true)).
+		WithField(entity.NewField().WithName(textField).WithDataType(entity.FieldTypeVarChar).WithMaxLength(2048).WithEnableAnalyzer(true).WithEnableMatch(true)).
+		WithField(entity.NewField().WithName(titleSparse).WithDataType(entity.FieldTypeSparseVector)).
+		WithField(entity.NewField().WithName(textSparse).WithDataType(entity.FieldTypeSparseVector)).
+		WithFunction(entity.NewFunction().WithName("title_bm25_func").WithType(entity.FunctionTypeBM25).WithInputFields(titleField).WithOutputFields(titleSparse)).
+		WithFunction(entity.NewFunction().WithName("text_bm25_func").WithType(entity.FunctionTypeBM25).WithInputFields(textField).WithOutputFields(textSparse))
+
+	idxOpts := []milvusclient.CreateIndexOption{
+		milvusclient.NewCreateIndexOption(collectionName, titleField, index.NewInvertedIndex()),
+		milvusclient.NewCreateIndexOption(collectionName, textField, index.NewInvertedIndex()),
+		milvusclient.NewCreateIndexOption(collectionName, titleSparse, index.NewSparseInvertedIndex(entity.BM25, 0.2)),
+		milvusclient.NewCreateIndexOption(collectionName, textSparse, index.NewSparseInvertedIndex(entity.BM25, 0.2)),
+	}
+
+	err = cli.CreateCollection(ctx, milvusclient.NewCreateCollectionOption(collectionName, schema).WithIndexOptions(idxOpts...))
+	if err != nil {
+		log.Fatal("failed to create collection: ", err.Error())
+	}
+
+	_, err = cli.Insert(ctx, milvusclient.NewColumnBasedInsertOption(collectionName).
+		WithVarcharColumn(titleField, []string{
+			"History of AI",
+			"Alan Turing Biography",
+			"Machine Learning Overview",
+		}).
+		WithVarcharColumn(textField, []string{
+			"Artificial intelligence was founded in 1956 by computer scientists.",
+			"Alan Turing proposed early concepts of AI and machine learning.",
+			"Machine learning is a subset of artificial intelligence.",
+		}))
+	if err != nil {
+		log.Fatal("failed to insert data: ", err.Error())
+	}
+
+	task, err := cli.LoadCollection(ctx, milvusclient.NewLoadCollectionOption(collectionName))
+	if err != nil {
+		log.Fatal("failed to load collection: ", err.Error())
+	}
+	_ = task.Await(ctx)
+
+	q := "artificial intelligence"
+	expr := "text_match(" + titleField + ", \"" + q + "\", minimum_should_match=2) OR text_match(" + textField + ", \"" + q + "\", minimum_should_match=2)"
+
+	boost := entity.NewFunction().
+		WithName("title_boost").
+		WithType(entity.FunctionTypeRerank).
+		WithParam("reranker", "boost").
+		WithParam("filter", "text_match("+titleField+", \""+q+"\", minimum_should_match=2)").
+		WithParam("weight", "2.0")
+
+	vectors := []entity.Vector{entity.Text(q)}
+	rs, err := cli.Search(ctx, milvusclient.NewSearchOption(collectionName, 5, vectors).
+		WithANNSField(textSparse).
+		WithFilter(expr).
+		WithOutputFields("id", titleField, textField).
+		WithFunctionReranker(boost))
+	if err != nil {
+		log.Fatal("failed to search: ", err.Error())
+	}
+
+	for _, r := range rs {
+		_ = r.ResultCount
+	}
+}

client/milvusclient/read_test.go

@@ -191,6 +191,46 @@ func (s *ReadSuite) TestSearch() {
 	})
 }
 
+// TestSearch_TextMatch tests the text match search functionality.
+// It tests the minimum_should_match parameter in the expression.
+func (s *ReadSuite) TestSearch_TextMatch() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	s.Run("min_should_match_in_expr", func() {
+		collectionName := fmt.Sprintf("coll_%s", s.randString(6))
+		s.setupCache(collectionName, s.schema)
+
+		s.mock.EXPECT().Search(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, sr *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
+			// ensure the expression contains minimum_should_match and both fields
+			s.Contains(sr.GetDsl(), "minimum_should_match=2")
+			s.Contains(sr.GetDsl(), "text_match(")
+			return &milvuspb.SearchResults{
+				Status: merr.Success(),
+				Results: &schemapb.SearchResultData{
+					NumQueries: 1,
+					TopK:       1,
+					FieldsData: []*schemapb.FieldData{
+						s.getInt64FieldData("ID", []int64{1}),
+					},
+					Ids:    &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
+					Scores: []float32{0.1},
+					Topks:  []int64{1},
+				},
+			}, nil
+		}).Once()
+
+		q := "artificial intelligence"
+		expr := "text_match(title, \"" + q + "\", minimum_should_match=2) OR text_match(document_text, \"" + q + "\", minimum_should_match=2)"
+		vectors := []entity.Vector{entity.Text(q)}
+		_, err := s.client.Search(ctx, NewSearchOption(collectionName, 5, vectors).
+			WithANNSField("text_sparse_vector").
+			WithFilter(expr).
+			WithOutputFields("ID"))
+		s.NoError(err)
+	})
+}
+
 func (s *ReadSuite) TestHybridSearch() {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -233,7 +273,7 @@ func (s *ReadSuite) TestHybridSearch() {
 			return rand.Float32()
 		}))).WithFilter("ID > 100"), NewAnnRequest("vector", 10, entity.FloatVector(lo.RepeatBy(128, func(_ int) float32 {
 			return rand.Float32()
-		})))).WithConsistencyLevel(entity.ClStrong).WithPartitons(partitionName).WithReranker(NewRRFReranker()).WithOutputFields("*"))
+		})))).WithConsistencyLevel(entity.ClStrong).WithPartitions(partitionName).WithReranker(NewRRFReranker()).WithOutputFields("*"))
 		s.NoError(err)
 	})
 

go.mod

@@ -65,7 +65,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.23.0
 	github.com/bits-and-blooms/bitset v1.10.0
 	github.com/bytedance/mockey v1.2.14
-	github.com/bytedance/sonic v1.13.2
+	github.com/bytedance/sonic v1.14.0
 	github.com/cenkalti/backoff/v4 v4.2.1
 	github.com/cockroachdb/redact v1.1.3
 	github.com/google/uuid v1.6.0
@@ -137,7 +137,7 @@ require (
 	github.com/benesch/cgosymbolizer v0.0.0-20190515212042-bec6fe6e597b // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/bytecodealliance/wasmtime-go v1.0.0 // indirect
-	github.com/bytedance/sonic/loader v0.2.4 // indirect
+	github.com/bytedance/sonic/loader v0.3.0 // indirect
 	github.com/campoy/embedmd v1.0.0 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/cilium/ebpf v0.11.0 // indirect

go.sum

@@ -207,11 +207,11 @@ github.com/bytecodealliance/wasmtime-go v1.0.0 h1:9u9gqaUiaJeN5IoD1L7egD8atOnTGy
 github.com/bytecodealliance/wasmtime-go v1.0.0/go.mod h1:jjlqQbWUfVSbehpErw3UoWFndBXRRMvfikYH6KsCwOg=
 github.com/bytedance/mockey v1.2.14 h1:KZaFgPdiUwW+jOWFieo3Lr7INM1P+6adO3hxZhDswY8=
 github.com/bytedance/mockey v1.2.14/go.mod h1:1BPHF9sol5R1ud/+0VEHGQq/+i2lN+GTsr3O2Q9IENY=
-github.com/bytedance/sonic v1.13.2 h1:8/H1FempDZqC4VqjptGo14QQlJx8VdZJegxs6wwfqpQ=
-github.com/bytedance/sonic v1.13.2/go.mod h1:o68xyaF9u2gvVBuGHPlUVCy+ZfmNNO5ETf1+KgkJhz4=
+github.com/bytedance/sonic v1.14.0 h1:/OfKt8HFw0kh2rj8N0F6C/qPGRESq0BbaNZgcNXXzQQ=
+github.com/bytedance/sonic v1.14.0/go.mod h1:WoEbx8WTcFJfzCe0hbmyTGrfjt8PzNEBdxlNUO24NhA=
 github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
-github.com/bytedance/sonic/loader v0.2.4 h1:ZWCw4stuXUsn1/+zQDqeE7JKP+QO47tz7QCNan80NzY=
-github.com/bytedance/sonic/loader v0.2.4/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI=
+github.com/bytedance/sonic/loader v0.3.0 h1:dskwH8edlzNMctoruo8FPTJDF3vLtDT0sXZwvZJyqeA=
+github.com/bytedance/sonic/loader v0.3.0/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI=
 github.com/campoy/embedmd v1.0.0 h1:V4kI2qTJJLf4J29RzI/MAt2c3Bl4dQSYPuflzwFH2hY=
 github.com/campoy/embedmd v1.0.0/go.mod h1:oxyr9RCiSXg0M3VJ3ks0UGfp98BpSSGr0kpiX3MzVl8=
 github.com/casbin/casbin/v2 v2.0.0/go.mod h1:YcPU1XXisHhLzuxH9coDNf2FbKpjGlbCg3n9yuLkIJQ=

internal/core/src/exec/expression/UnaryExpr.cpp

@@ -1792,10 +1792,18 @@ PhyUnaryRangeFilterExpr::ExecTextMatch() {
         }
     }
 
-    auto func = [op_type, slop](Index* index,
-                                const std::string& query) -> TargetBitmap {
+    uint32_t min_should_match = 1;  // default value
+    if (op_type == proto::plan::OpType::TextMatch &&
+        expr_->extra_values_.size() > 0) {
+        // min_should_match is stored in the first extra value
+        min_should_match = static_cast<uint32_t>(
+            GetValueFromProto<int64_t>(expr_->extra_values_[0]));
+    }
+
+    auto func = [op_type, slop, min_should_match](
+                    Index* index, const std::string& query) -> TargetBitmap {
         if (op_type == proto::plan::OpType::TextMatch) {
-            return index->MatchQuery(query);
+            return index->MatchQuery(query, min_should_match);
         } else if (op_type == proto::plan::OpType::PhraseMatch) {
             return index->PhraseMatchQuery(query, slop);
         } else {

internal/core/src/index/TextMatchIndex.cpp

@@ -153,12 +153,12 @@ TextMatchIndex::Load(const Config& config) {
         std::vector<std::string> file;
         file.push_back(*it);
         files_value.erase(it);
-        auto index_datas =
+        auto index_data =
             mem_file_manager_->LoadIndexToMemory(file, load_priority);
         BinarySet binary_set;
-        AssembleIndexDatas(index_datas, binary_set);
-        // clear index_datas to free memory early
-        index_datas.clear();
+        AssembleIndexDatas(index_data, binary_set);
+        // clear index_data to free memory early
+        index_data.clear();
         auto index_valid_data = binary_set.GetByName("index_null_offset");
         null_offset_.resize((size_t)index_valid_data->size / sizeof(size_t));
         memcpy(null_offset_.data(),
@@ -226,18 +226,18 @@ TextMatchIndex::AddTextsGrowing(size_t n,
 // schema_ may not be initialized so we need this `nullable` parameter
 void
 TextMatchIndex::BuildIndexFromFieldData(
-    const std::vector<FieldDataPtr>& field_datas, bool nullable) {
+    const std::vector<FieldDataPtr>& field_data, bool nullable) {
     int64_t offset = 0;
     if (nullable) {
         int64_t total = 0;
-        for (const auto& data : field_datas) {
+        for (const auto& data : field_data) {
             total += data->get_null_count();
         }
         {
             std::unique_lock<folly::SharedMutex> lock(mutex_);
             null_offset_.reserve(total);
         }
-        for (const auto& data : field_datas) {
+        for (const auto& data : field_data) {
             auto n = data->get_num_rows();
             for (int i = 0; i < n; i++) {
                 if (!data->is_valid(i)) {
@@ -251,7 +251,7 @@ TextMatchIndex::BuildIndexFromFieldData(
             }
         }
     } else {
-        for (const auto& data : field_datas) {
+        for (const auto& data : field_data) {
             auto n = data->get_num_rows();
             wrapper_->add_data(
                 static_cast<const std::string*>(data->Data()), n, offset);
@@ -302,17 +302,19 @@ TextMatchIndex::RegisterTokenizer(const char* tokenizer_name,
 }
 
 TargetBitmap
-TextMatchIndex::MatchQuery(const std::string& query) {
+TextMatchIndex::MatchQuery(const std::string& query,
+                           uint32_t min_should_match) {
+    tracer::AutoSpan span("TextMatchIndex::MatchQuery", tracer::GetRootSpan());
     if (shouldTriggerCommit()) {
         Commit();
         Reload();
     }
 
     TargetBitmap bitset{static_cast<size_t>(Count())};
-    // The count opeartion of tantivy may be get older cnt if the index is committed with new tantivy segment.
+    // The count operation of tantivy may be get older cnt if the index is committed with new tantivy segment.
     // So we cannot use the count operation to get the total count for bitmap.
     // Just use the maximum offset of hits to get the total count for bitmap here.
-    wrapper_->match_query(query, &bitset);
+    wrapper_->match_query(query, min_should_match, &bitset);
     return bitset;
 }
 
@@ -324,7 +326,7 @@ TextMatchIndex::PhraseMatchQuery(const std::string& query, uint32_t slop) {
     }
 
     TargetBitmap bitset{static_cast<size_t>(Count())};
-    // The count opeartion of tantivy may be get older cnt if the index is committed with new tantivy segment.
+    // The count operation of tantivy may be get older cnt if the index is committed with new tantivy segment.
     // So we cannot use the count operation to get the total count for bitmap.
     // Just use the maximum offset of hits to get the total count for bitmap here.
     wrapper_->phrase_match_query(query, slop, &bitset);

internal/core/src/index/TextMatchIndex.h

@@ -63,7 +63,7 @@ class TextMatchIndex : public InvertedIndexTantivy<std::string> {
                     int64_t offset_begin);
 
     void
-    BuildIndexFromFieldData(const std::vector<FieldDataPtr>& field_datas,
+    BuildIndexFromFieldData(const std::vector<FieldDataPtr>& field_data,
                             bool nullable);
 
     void
@@ -83,7 +83,7 @@ class TextMatchIndex : public InvertedIndexTantivy<std::string> {
     RegisterTokenizer(const char* tokenizer_name, const char* analyzer_params);
 
     TargetBitmap
-    MatchQuery(const std::string& query);
+    MatchQuery(const std::string& query, uint32_t min_should_match);
 
     TargetBitmap
     PhraseMatchQuery(const std::string& query, uint32_t slop);

internal/core/src/index/TextMatchIndexTest.cpp

@@ -164,7 +164,7 @@ TEST(TextMatch, Index) {
     index->Reload();
 
     {
-        auto res = index->MatchQuery("football");
+        auto res = index->MatchQuery("football", 1);
         ASSERT_EQ(res.size(), 3);
         ASSERT_TRUE(res[0]);
         ASSERT_FALSE(res[1]);
@@ -177,11 +177,16 @@ TEST(TextMatch, Index) {
         ASSERT_TRUE(res2[0]);
         ASSERT_FALSE(res2[1]);
         ASSERT_TRUE(res2[2]);
-        res = index->MatchQuery("nothing");
+        res = index->MatchQuery("nothing", 1);
         ASSERT_EQ(res.size(), 3);
         ASSERT_FALSE(res[0]);
         ASSERT_FALSE(res[1]);
         ASSERT_FALSE(res[2]);
+        auto res3 = index->MatchQuery("football pingpang cricket", 2);
+        ASSERT_EQ(res3.size(), 3);
+        ASSERT_TRUE(res3[0]);
+        ASSERT_FALSE(res3[1]);
+        ASSERT_FALSE(res3[2]);
     }
 
     {