Merge pull request #18 from filipecosta90/spellcheck

[add] added support for FT.SPELLCHECK

Author: filipecosta90

redisearch/client.go

@@ -346,6 +346,49 @@ func (i *Client) Search(q *Query) (docs []Document, total int, err error) {
 	return
 }
 
+// SpellCheck performs spelling correction on a query, returning suggestions for misspelled terms,
+// the total number of results, or an error if something went wrong
+func (i *Client) SpellCheck(q *Query, s *SpellCheckOptions) (suggs []MisspelledTerm, total int, err error) {
+	conn := i.pool.Get()
+	defer conn.Close()
+
+	args := redis.Args{i.name}
+	args = append(args, q.serialize()...)
+	args = append(args, s.serialize()...)
+
+	res, err := redis.Values(conn.Do("FT.SPELLCHECK", args...))
+	if err != nil {
+		return
+	}
+	total = 0
+	suggs = make([]MisspelledTerm, 0 )
+
+	// Each misspelled term, in turn, is a 3-element array consisting of
+	// - the constant string "TERM" ( 3-element position 0 -- we dont use it )
+	// - the term itself ( 3-element position 1 )
+	// - an array of suggestions for spelling corrections ( 3-element position 2 )
+	termIdx := 1
+	suggIdx := 2
+	for i := 0; i < len(res); i ++ {
+		var termArray []interface{} = nil
+		termArray, err = redis.Values(res[i], nil)
+		if err != nil {
+			return
+		}
+
+		if d, e := loadMisspelledTerm(termArray, termIdx, suggIdx); e == nil {
+			suggs = append(suggs, d)
+			if d.Len() > 0 {
+				total++
+			}
+		} else {
+			log.Print("Error parsing misspelled suggestion: ", e)
+		}
+	}
+
+	return
+}
+
 // Aggregate
 func (i *Client) Aggregate( q *AggregateQuery ) ( aggregateReply [][]string, total int, err error) {
 	conn := i.pool.Get()
@@ -389,7 +432,6 @@ func (i *Client) Aggregate( q *AggregateQuery ) ( aggregateReply [][]string, tot
 	return
 }
 
-
 // Explain Return a textual string explaining the query
 func (i *Client) Explain(q *Query) (string, error) {
 	conn := i.pool.Get()

redisearch/redisearch_test.go

@@ -383,6 +383,47 @@ func TestDelete(t *testing.T) {
 	assert.Equal(t, uint64(0), info.DocCount)
 }
 
+func TestSpellCheck(t *testing.T) {
+	c := createClient("testung")
+	countries := []string{"Spain", "Israel", "Portugal", "France", "England", "Angola"}
+	sc := redisearch.NewSchema(redisearch.DefaultOptions).
+		AddField(redisearch.NewTextField("country"))
+	c.Drop()
+	assert.Nil(t, c.CreateIndex(sc))
+
+	docs := make([]redisearch.Document, len(countries))
+
+	for i := 0; i < len(countries); i++ {
+		docs[i] = redisearch.NewDocument(fmt.Sprintf("doc%d", i), 1).Set("country", countries[i])
+	}
+
+	assert.Nil(t, c.Index(docs...))
+	query := redisearch.NewQuery("Anla Portuga" )
+	opts := redisearch.NewSpellCheckOptions(2 )
+	sugs, total, err := c.SpellCheck(query,opts )
+	assert.Nil(t, err)
+	assert.Equal(t, 2, len(sugs))
+	assert.Equal(t, 2, total)
+
+	// query that return the MisspelledTerm but with an empty list of suggestions
+	// 1) 1) "TERM"
+	//   2) "an"
+	//   3) (empty list or set)
+	queryEmpty := redisearch.NewQuery("An" )
+	sugs, total, err = c.SpellCheck(queryEmpty,opts )
+	assert.Nil(t, err)
+	assert.Equal(t, 1, len(sugs))
+	assert.Equal(t, 0, total)
+
+	// same query but now with a distance of 4
+	opts.SetDistance(4)
+	sugs, total, err = c.SpellCheck(queryEmpty,opts )
+	assert.Nil(t, err)
+	assert.Equal(t, 1, len(sugs))
+	assert.Equal(t, 1, total)
+
+}
+
 func ExampleClient() {
 
 	// Create a client. By default a client is schemaless

redisearch/spellcheck.go

@@ -0,0 +1,137 @@
+package redisearch
+
+import (
+	"fmt"
+	"github.com/gomodule/redigo/redis"
+	"sort"
+)
+
+// SpellCheckOptions are options which are passed when performing spelling correction on a query
+type SpellCheckOptions struct {
+	Distance       int
+	ExclusionDicts []string
+	InclusionDicts []string
+}
+
+func NewSpellCheckOptionsDefaults() *SpellCheckOptions {
+	return &SpellCheckOptions{
+		Distance:       1,
+		ExclusionDicts: make([]string, 0),
+		InclusionDicts: make([]string, 0),
+	}
+}
+
+func NewSpellCheckOptions(distance int) *SpellCheckOptions {
+	return &SpellCheckOptions{
+		Distance:       distance,
+		ExclusionDicts: make([]string, 0),
+		InclusionDicts: make([]string, 0),
+	}
+}
+
+// SetDistance Sets the the maximal Levenshtein distance for spelling suggestions (default: 1, max: 4)
+func (s *SpellCheckOptions) SetDistance(distance int) (*SpellCheckOptions, error) {
+	if distance < 1 || distance > 4 {
+		return s, fmt.Errorf("The maximal Levenshtein distance for spelling suggestions should be between [1,4]. Got %d", distance)
+	} else {
+		s.Distance = distance
+	}
+	return s, nil
+}
+
+// AddExclusionDict adds a custom dictionary named {dictname} to the exclusion list
+func (s *SpellCheckOptions) AddExclusionDict(dictname string) *SpellCheckOptions {
+	s.ExclusionDicts = append(s.ExclusionDicts, dictname)
+	return s
+}
+
+// AddInclusionDict adds a custom dictionary named {dictname} to the inclusion list
+func (s *SpellCheckOptions) AddInclusionDict(dictname string) *SpellCheckOptions {
+	s.InclusionDicts = append(s.InclusionDicts, dictname)
+	return s
+}
+
+func (s SpellCheckOptions) serialize() redis.Args {
+	args := redis.Args{}
+	if s.Distance > 1 {
+		args = args.Add("DISTANCE").Add(s.Distance)
+	}
+	for _, exclusion := range s.ExclusionDicts {
+		args = args.Add("TERMS").Add("EXCLUDE").Add(exclusion)
+	}
+	for _, inclusion := range s.InclusionDicts {
+		args = args.Add("TERMS").Add("INCLUDE").Add(inclusion)
+	}
+	return args
+}
+
+// MisspelledSuggestion is a single suggestion from the spelling corrections
+type MisspelledSuggestion struct {
+	Suggestion string
+	Score      float32
+}
+
+// NewMisspelledSuggestion creates a MisspelledSuggestion with the specific term and score
+func NewMisspelledSuggestion(term string, score float32) MisspelledSuggestion {
+	return MisspelledSuggestion{
+		Suggestion: term,
+		Score:      score,
+	}
+}
+
+// MisspelledTerm contains the misspelled term and a sortable list of suggestions returned from an engine
+type MisspelledTerm struct {
+	Term string
+	// MisspelledSuggestionList is a sortable list of suggestions returned from an engine
+	MisspelledSuggestionList []MisspelledSuggestion
+}
+
+func NewMisspelledTerm(term string) MisspelledTerm {
+	return MisspelledTerm{
+		Term:                     term,
+		MisspelledSuggestionList: make([]MisspelledSuggestion, 0),
+	}
+}
+
+func (l MisspelledTerm) Len() int { return len(l.MisspelledSuggestionList) }
+func (l MisspelledTerm) Swap(i, j int) {
+	l.MisspelledSuggestionList[i], l.MisspelledSuggestionList[j] = l.MisspelledSuggestionList[j], l.MisspelledSuggestionList[i]
+}
+func (l MisspelledTerm) Less(i, j int) bool {
+	return l.MisspelledSuggestionList[i].Score > l.MisspelledSuggestionList[j].Score
+} //reverse sorting
+
+// Sort the SuggestionList
+func (l MisspelledTerm) Sort() {
+	sort.Sort(l)
+}
+
+// convert the result from a redis spelling correction on a query to a proper MisspelledTerm object
+func loadMisspelledTerm(arr []interface{}, termIdx, suggIdx int) (missT MisspelledTerm, err error) {
+	term, err := redis.String(arr[termIdx], err)
+	if err != nil {
+		return MisspelledTerm{}, fmt.Errorf("Could not parse term: %s", err)
+	}
+	missT = NewMisspelledTerm(term)
+	lst, err := redis.Values(arr[suggIdx], err)
+	if err != nil {
+		return MisspelledTerm{}, fmt.Errorf("Could not get the array of suggestions for spelling corrections on term %s. Error: %s", term, err)
+	}
+	for i := 0; i < len(lst); i++ {
+		innerLst, err := redis.Values(lst[i], err)
+		if err != nil {
+			return MisspelledTerm{}, fmt.Errorf("Could not get the inner array of suggestions for spelling corrections on term %s. Error: %s", term, err)
+		}
+		score, err := redis.Float64(innerLst[0], err)
+		if err != nil {
+			return MisspelledTerm{}, fmt.Errorf("Could not parse score: %s", err)
+		}
+		suggestion, err := redis.String(innerLst[1], err)
+		if err != nil {
+			return MisspelledTerm{}, fmt.Errorf("Could not parse suggestion: %s", err)
+		}
+		missT.MisspelledSuggestionList = append(missT.MisspelledSuggestionList, NewMisspelledSuggestion(suggestion, float32(score)))
+	}
+
+	return missT, nil
+}