Skip to content

Commit

Permalink
Merge pull request #509 from redis/asftsearch-scores
Browse files Browse the repository at this point in the history 
feat: parse FT.SEARCH WITHSCORES in AsFtSearch()
  • Loading branch information
@rueian
rueian authored Mar 22, 2024
2 parents f7e1abb + 0bb92a3 commit 04d4621
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 20 deletions.
44 changes: 33 additions & 11 deletions message.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1014,8 +1014,9 @@ func (m *RedisMessage) AsZMPop() (kvs KeyZScores, err error) {
}

type FtSearchDoc struct {
Doc map[string]string
Key string
Doc map[string]string
Key string
Score float64
}

func (m *RedisMessage) AsFtSearch() (total int64, docs []FtSearchDoc, err error) {
Expand All @@ -1037,6 +1038,8 @@ func (m *RedisMessage) AsFtSearch() (total int64, docs []FtSearchDoc, err error)
docs[d].Key = record.values[j+1].string
case "extra_attributes":
docs[d].Doc, _ = record.values[j+1].AsStrMap()
case "score":
docs[d].Score, _ = strconv.ParseFloat(record.values[j+1].string, 64)
}
}
}
Expand All @@ -1051,17 +1054,36 @@ func (m *RedisMessage) AsFtSearch() (total int64, docs []FtSearchDoc, err error)
}
if len(m.values) > 0 {
total = m.values[0].integer
if len(m.values) > 2 && m.values[2].string == "" {
docs = make([]FtSearchDoc, 0, (len(m.values)-1)/2)
for i := 1; i < len(m.values); i += 2 {
doc, _ := m.values[i+1].AsStrMap()
docs = append(docs, FtSearchDoc{Doc: doc, Key: m.values[i].string})
wscore := false
wattrs := false
offset := 1
if len(m.values) > 2 {
if m.values[2].string == "" {
wattrs = true
offset++
} else {
_, err1 := strconv.ParseFloat(m.values[1].string, 64)
_, err2 := strconv.ParseFloat(m.values[2].string, 64)
wscore = err1 != nil && err2 == nil
offset++
}
} else {
docs = make([]FtSearchDoc, 0, len(m.values)-1)
for i := 1; i < len(m.values); i++ {
docs = append(docs, FtSearchDoc{Doc: nil, Key: m.values[i].string})
}
if len(m.values) > 3 && m.values[3].string == "" {
wattrs = true
offset++
}
docs = make([]FtSearchDoc, 0, (len(m.values)-1)/offset)
for i := 1; i < len(m.values); i++ {
doc := FtSearchDoc{Key: m.values[i].string}
if wscore {
i++
doc.Score, _ = strconv.ParseFloat(m.values[i].string, 64)
}
if wattrs {
i++
doc.Doc, _ = m.values[i].AsStrMap()
}
docs = append(docs, doc)
}
return
}
Expand Down
66 changes: 57 additions & 9 deletions message_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -610,55 +610,97 @@ func TestRedisResult(t *testing.T) {
}
if n, ret, _ := (RedisResult{val: RedisMessage{typ: '*', values: []RedisMessage{
{typ: ':', integer: 3},
{typ: '+', string: "1"},
{typ: '+', string: "a"},
{typ: '*', values: []RedisMessage{
{typ: '+', string: "k1"},
{typ: '+', string: "v1"},
{typ: '+', string: "kk"},
{typ: '+', string: "vv"},
}},
{typ: '+', string: "2"},
{typ: '+', string: "b"},
{typ: '*', values: []RedisMessage{
{typ: '+', string: "k2"},
{typ: '+', string: "v2"},
{typ: '+', string: "kk"},
{typ: '+', string: "vv"},
}},
}}}).AsFtSearch(); n != 3 || !reflect.DeepEqual([]FtSearchDoc{
{Key: "1", Doc: map[string]string{"k1": "v1", "kk": "vv"}},
{Key: "2", Doc: map[string]string{"k2": "v2", "kk": "vv"}},
{Key: "a", Doc: map[string]string{"k1": "v1", "kk": "vv"}},
{Key: "b", Doc: map[string]string{"k2": "v2", "kk": "vv"}},
}, ret) {
t.Fatal("AsFtSearch not get value as expected")
}
if n, ret, _ := (RedisResult{val: RedisMessage{typ: '*', values: []RedisMessage{
{typ: ':', integer: 3},
{typ: '+', string: "a"},
{typ: '+', string: "1"},
{typ: '*', values: []RedisMessage{
{typ: '+', string: "k1"},
{typ: '+', string: "v1"},
}},
{typ: '+', string: "b"},
{typ: '+', string: "2"},
{typ: '*', values: []RedisMessage{
{typ: '+', string: "k2"},
{typ: '+', string: "v2"},
}},
}}}).AsFtSearch(); n != 3 || !reflect.DeepEqual([]FtSearchDoc{
{Key: "a", Doc: map[string]string{"k1": "v1"}, Score: 1},
{Key: "b", Doc: map[string]string{"k2": "v2"}, Score: 2},
}, ret) {
t.Fatal("AsFtSearch not get value as expected")
}
if n, ret, _ := (RedisResult{val: RedisMessage{typ: '*', values: []RedisMessage{
{typ: ':', integer: 3},
{typ: '+', string: "a"},
{typ: '*', values: []RedisMessage{
{typ: '+', string: "k1"},
{typ: '+', string: "v1"},
{typ: '+', string: "kk"},
{typ: '+', string: "vv"},
}},
}}}).AsFtSearch(); n != 3 || !reflect.DeepEqual([]FtSearchDoc{
{Key: "1", Doc: map[string]string{"k1": "v1", "kk": "vv"}},
{Key: "a", Doc: map[string]string{"k1": "v1", "kk": "vv"}},
}, ret) {
t.Fatal("AsFtSearch not get value as expected")
}
if n, ret, _ := (RedisResult{val: RedisMessage{typ: '*', values: []RedisMessage{
{typ: ':', integer: 3},
{typ: '+', string: "a"},
{typ: '+', string: "b"},
}}}).AsFtSearch(); n != 3 || !reflect.DeepEqual([]FtSearchDoc{
{Key: "a", Doc: nil},
{Key: "b", Doc: nil},
}, ret) {
t.Fatal("AsFtSearch not get value as expected")
}
if n, ret, _ := (RedisResult{val: RedisMessage{typ: '*', values: []RedisMessage{
{typ: ':', integer: 3},
{typ: '+', string: "a"},
{typ: '+', string: "1"},
{typ: '+', string: "b"},
{typ: '+', string: "2"},
}}}).AsFtSearch(); n != 3 || !reflect.DeepEqual([]FtSearchDoc{
{Key: "1", Doc: nil},
{Key: "2", Doc: nil},
{Key: "a", Doc: nil, Score: 1},
{Key: "b", Doc: nil, Score: 2},
}, ret) {
t.Fatal("AsFtSearch not get value as expected")
}
if n, ret, _ := (RedisResult{val: RedisMessage{typ: '*', values: []RedisMessage{
{typ: ':', integer: 3},
{typ: '+', string: "1"},
{typ: '+', string: "2"},
}}}).AsFtSearch(); n != 3 || !reflect.DeepEqual([]FtSearchDoc{
{Key: "1", Doc: nil},
{Key: "2", Doc: nil},
}, ret) {
t.Fatal("AsFtSearch not get value as expected")
}
if n, ret, _ := (RedisResult{val: RedisMessage{typ: '*', values: []RedisMessage{
{typ: ':', integer: 3},
{typ: '+', string: "a"},
}}}).AsFtSearch(); n != 3 || !reflect.DeepEqual([]FtSearchDoc{
{Key: "a", Doc: nil},
}, ret) {
t.Fatal("AsFtSearch not get value as expected")
}
Expand All @@ -683,6 +725,8 @@ func TestRedisResult(t *testing.T) {
{typ: '+', string: "$"},
{typ: '+', string: "1"},
}},
{typ: '+', string: "score"},
{typ: ',', string: "1"},
}},
{typ: '%', values: []RedisMessage{
{typ: '+', string: "id"},
Expand All @@ -692,13 +736,15 @@ func TestRedisResult(t *testing.T) {
{typ: '+', string: "$"},
{typ: '+', string: "2"},
}},
{typ: '+', string: "score"},
{typ: ',', string: "2"},
}},
}},
{typ: '+', string: "error"},
{typ: '*', values: []RedisMessage{}},
}}}).AsFtSearch(); n != 3 || !reflect.DeepEqual([]FtSearchDoc{
{Key: "1", Doc: map[string]string{"$": "1"}},
{Key: "2", Doc: map[string]string{"$": "2"}},
{Key: "1", Doc: map[string]string{"$": "1"}, Score: 1},
{Key: "2", Doc: map[string]string{"$": "2"}, Score: 2},
}, ret) {
t.Fatal("AsFtSearch not get value as expected")
}
Expand All @@ -715,6 +761,8 @@ func TestRedisResult(t *testing.T) {
{typ: '+', string: "$"},
{typ: '+', string: "1"},
}},
{typ: '+', string: "score"},
{typ: ',', string: "1"},
}},
}},
{typ: '+', string: "error"},
Expand Down

0 comments on commit 04d4621

Please sign in to comment.