Add Bleve for fulltext search

internal/fulltext/bleve.go

@@ -1,7 +1,26 @@
+// Copyright 2022 The Matrix.org Foundation C.I.C.
+//
+// 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.
+
 package fulltext
 
 import (
+	"strings"
+	"time"
+
 	"github.com/blevesearch/bleve/v2"
+	"github.com/blevesearch/bleve/v2/analysis/lang/en"
+	"github.com/blevesearch/bleve/v2/search/query"
 )
 
 // Search contains all existing bleve.Index
@@ -11,10 +30,10 @@ type Search struct {
 
 // IndexElement describes the layout of an element to index
 type IndexElement struct {
-	EventID string
-	Type    string
-	RoomID  string
-	Content string
+	EventID string    `json:"event_id,omitempty"`
+	RoomID  string    `json:"room_id,omitempty"`
+	Content string    `json:"content,omitempty"`
+	Time    time.Time `json:"timestamp,omitempty"`
 }
 
 // New opens a new/existing fulltext index
@@ -38,16 +57,50 @@ func (f *Search) Index(e IndexElement) error {
 	return f.MessageIndex.Index(e.EventID, e)
 }
 
+// BatchIndex indexes the given elements
+func (f *Search) BatchIndex(elements []IndexElement) error {
+	batch := f.MessageIndex.NewBatch()
+
+	for _, element := range elements {
+		err := batch.Index(element.EventID, element)
+		if err != nil {
+			return err
+		}
+	}
+	return f.MessageIndex.Batch(batch)
+}
+
 // Delete deletes an indexed element by the eventID
 func (f *Search) Delete(eventID string) error {
 	return f.MessageIndex.Delete(eventID)
 }
 
 // Search searches the index given a search term
-func (f *Search) Search(term string) (*bleve.SearchResult, error) {
-	qry := bleve.NewQueryStringQuery(term)
-	search := bleve.NewSearchRequest(qry)
-	return f.MessageIndex.Search(search)
+func (f *Search) Search(term string, roomIDs []string, limit, from int, orderByTime bool) (*bleve.SearchResult, error) {
+	terms := strings.Split(term, " ")
+
+	qry := bleve.NewConjunctionQuery()
+	for _, t := range terms {
+		qry.AddQuery(bleve.NewQueryStringQuery(t))
+	}
+
+	for _, roomID := range roomIDs {
+		roomSearch := bleve.NewMatchQuery(roomID)
+		roomSearch.SetField("room_id")
+		roomSearch.SetOperator(query.MatchQueryOperatorAnd)
+		qry.AddQuery(roomSearch)
+	}
+
+	s := bleve.NewSearchRequest(qry)
+	s.Size = limit
+	s.From = from
+
+	s.SortBy([]string{"_score"})
+	if orderByTime {
+		s.SortBy([]string{"-timestamp"})
+	}
+
+	return f.MessageIndex.Search(s)
 }
 
 func openIndex(path string) (bleve.Index, error) {
@@ -55,9 +108,47 @@ func openIndex(path string) (bleve.Index, error) {
 		return index, nil
 	}
 
-	index, err := bleve.New(path, bleve.NewIndexMapping())
+	enFieldMapping := bleve.NewTextFieldMapping()
+	enFieldMapping.Analyzer = en.AnalyzerName
+
+	eventMapping := bleve.NewDocumentMapping()
+
+	eventMapping.AddFieldMappingsAt("content", enFieldMapping)
+	eventMapping.AddFieldMappingsAt("room_id", bleve.NewTextFieldMapping())
+
+	idMapping := bleve.NewTextFieldMapping()
+	idMapping.IncludeInAll = false
+	idMapping.Index = false
+	idMapping.IncludeTermVectors = false
+	idMapping.SkipFreqNorm = true
+	eventMapping.AddFieldMappingsAt("event_id", idMapping)
+
+	mapping := bleve.NewIndexMapping()
+	mapping.AddDocumentMapping("event", eventMapping)
+	mapping.DefaultType = "event"
+	mapping.TypeField = "type"
+	mapping.DefaultAnalyzer = "en"
+
+	index, err := bleve.New(path, mapping)
 	if err != nil {
 		return nil, err
 	}
 	return index, nil
 }
+
+type IndexElements []IndexElement
+
+// Len implements sort.Interface
+func (ie IndexElements) Len() int {
+	return len(ie)
+}
+
+// Less implements sort.Interface
+func (ie IndexElements) Less(i, j int) bool {
+	return ie[i].Time.After(ie[j].Time)
+}
+
+// Swap implements sort.Interface
+func (ie IndexElements) Swap(i, j int) {
+	ie[i], ie[j] = ie[j], ie[i]
+}

internal/fulltext/bleve_test.go

@@ -1,15 +1,32 @@
-package fulltext
+// Copyright 2022 The Matrix.org Foundation C.I.C.
+//
+// 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.
+
+package fulltext_test
 
 import (
+	"sort"
 	"testing"
+	"time"
 
+	"github.com/matrix-org/dendrite/internal/fulltext"
 	"github.com/matrix-org/util"
 )
 
 func TestSearch(t *testing.T) {
 	// create new index
 	dataDir := t.TempDir()
-	fts, err := New(dataDir)
+	fts, err := fulltext.New(dataDir)
 	if err != nil {
 		t.Fatal("failed to open fulltext index:", err)
 	}
@@ -18,7 +35,7 @@ func TestSearch(t *testing.T) {
 	}
 
 	// open existing index
-	fts, err = New(dataDir)
+	fts, err = fulltext.New(dataDir)
 	if err != nil {
 		t.Fatal("failed to open fulltext index:", err)
 	}
@@ -28,11 +45,12 @@ func TestSearch(t *testing.T) {
 	}
 
 	// add some data
-	e := IndexElement{
+	roomID := util.RandomString(8)
+	e := fulltext.IndexElement{
 		EventID: util.RandomString(16),
-		Type:    "m.room.message",
-		RoomID:  util.RandomString(8),
+		RoomID:  roomID,
 		Content: "lorem ipsum",
+		Time:    time.Now(),
 	}
 
 	if err = fts.Index(e); err != nil {
@@ -40,11 +58,11 @@ func TestSearch(t *testing.T) {
 	}
 
 	eventID := util.RandomString(16)
-	e = IndexElement{
+	e = fulltext.IndexElement{
 		EventID: eventID,
-		Type:    "m.room.message",
-		RoomID:  util.RandomString(8),
+		RoomID:  roomID,
 		Content: "lorem ipsum",
+		Time:    time.Now(),
 	}
 
 	if err = fts.Index(e); err != nil {
@@ -52,7 +70,7 @@ func TestSearch(t *testing.T) {
 	}
 
 	// search data
-	res, err := fts.Search("lorem")
+	res, err := fts.Search("lorem", nil, 10, 0, false)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -65,11 +83,51 @@ func TestSearch(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	res, err = fts.Search("lorem")
+	// create some more random data
+	var batchItems []fulltext.IndexElement
+
+	wantRoomID := util.RandomString(8)
+	for i := 0; i < 30; i++ {
+		eventID = util.RandomString(16)
+		e = fulltext.IndexElement{
+			EventID: eventID,
+			RoomID:  wantRoomID,
+			Content: "lorem ipsum",
+			Time:    time.Now(),
+		}
+		batchItems = append(batchItems, e)
+	}
+
+	// Index the data
+	if err = fts.BatchIndex(batchItems); err != nil {
+		t.Fatal("failed to batch index")
+	}
+
+	// search for lorem, but only in a given room
+	searchRooms := []string{roomID}
+	res, err = fts.Search("lorem", searchRooms, 10, 0, false)
 	if err != nil {
 		t.Fatal(err)
 	}
 	if res.Total != 1 {
-		t.Fatalf("expected %d results, got %d", 2, res.Total)
+		t.Fatalf("expected %d results, got %d", 1, res.Total)
 	}
+
+	// can get sorted results
+	res, err = fts.Search("lorem", []string{wantRoomID}, 10, 0, true)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if res.Hits[0].ID != eventID {
+		t.Fatalf("expected %s to be first, got %s", eventID, res.Hits[0].ID)
+	}
+
+	sort.Sort(fulltext.IndexElements(batchItems))
+	if eventID != batchItems[0].EventID {
+		t.Fatalf("expected %s to be first, got %s", eventID, batchItems[0].EventID)
+	}
+
+	// test back pagination
+
 }