mirror of
https://github.com/redis/go-redis.git
synced 2025-06-16 23:40:55 +03:00
4e22885ca1
* feat: support vectorset
* fix: char encoding error
* use `any` instread of `interface{}`
* update vectorset API
Signed-off-by: fukua95 <fukua95@gmail.com>
* refact: MapStringFloat64Cmd -> VectorInfoSliceCmd
Signed-off-by: fukua95 <fukua95@gmail.com>
* update:
* the type of vector attribute: string -> VectorAttributeMarshaller
* Add a new API VRemAttr
* mark the APIs are experimental
Signed-off-by: fukua95 <fukua95@gmail.com>
* trigger CI again
Signed-off-by: fukua95 <fukua95@gmail.com>
* rename a API: VRemAttr -> VClearAttributes
Signed-off-by: fukua95 <fukua95@gmail.com>
* add test
Signed-off-by: fukua95 <fukua95@gmail.com>
* feat(vectorset): improve VSetAttr API and add comprehensive test suite
- Simplify VSetAttr to accept interface{} with automatic JSON marshalling
- Remove VectorAttributeMarshaller interface for simpler API
- Add comprehensive unit tests for all vectorset commands
---------
Signed-off-by: fukua95 <fukua95@gmail.com>
Co-authored-by: Nedyalko Dyakov <nedyalko.dyakov@gmail.com>
327 lines
8.7 KiB
Go
327 lines
8.7 KiB
Go
package redis_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math/rand"
|
|
"time"
|
|
|
|
. "github.com/bsm/ginkgo/v2"
|
|
. "github.com/bsm/gomega"
|
|
"github.com/redis/go-redis/v9"
|
|
"github.com/redis/go-redis/v9/internal/proto"
|
|
)
|
|
|
|
func expectNil(err error) {
|
|
Expect(err).NotTo(HaveOccurred())
|
|
}
|
|
|
|
func expectTrue(t bool) {
|
|
expectEqual(t, true)
|
|
}
|
|
|
|
func expectEqual[T any, U any](a T, b U) {
|
|
Expect(a).To(BeEquivalentTo(b))
|
|
}
|
|
|
|
func generateRandomVector(dim int) redis.VectorValues {
|
|
rand.Seed(time.Now().UnixNano())
|
|
v := make([]float64, dim)
|
|
for i := range v {
|
|
v[i] = float64(rand.Intn(1000)) + rand.Float64()
|
|
}
|
|
return redis.VectorValues{Val: v}
|
|
}
|
|
|
|
var _ = Describe("Redis VectorSet commands", Label("vectorset"), func() {
|
|
ctx := context.TODO()
|
|
|
|
setupRedisClient := func(protocolVersion int) *redis.Client {
|
|
return redis.NewClient(&redis.Options{
|
|
Addr: "localhost:6379",
|
|
DB: 0,
|
|
Protocol: protocolVersion,
|
|
UnstableResp3: true,
|
|
})
|
|
}
|
|
|
|
protocols := []int{2, 3}
|
|
for _, protocol := range protocols {
|
|
protocol := protocol
|
|
|
|
Context(fmt.Sprintf("with protocol version %d", protocol), func() {
|
|
var client *redis.Client
|
|
|
|
BeforeEach(func() {
|
|
client = setupRedisClient(protocol)
|
|
Expect(client.FlushAll(ctx).Err()).NotTo(HaveOccurred())
|
|
})
|
|
|
|
AfterEach(func() {
|
|
if client != nil {
|
|
client.FlushDB(ctx)
|
|
client.Close()
|
|
}
|
|
})
|
|
|
|
It("basic", func() {
|
|
SkipBeforeRedisVersion(8.0, "Redis 8.0 introduces support for VectorSet")
|
|
vecName := "basic"
|
|
val := &redis.VectorValues{
|
|
Val: []float64{1.5, 2.4, 3.3, 4.2},
|
|
}
|
|
ok, err := client.VAdd(ctx, vecName, "k1", val).Result()
|
|
expectNil(err)
|
|
expectTrue(ok)
|
|
|
|
fp32 := "\x8f\xc2\xf9\x3e\xcb\xbe\xe9\xbe\xb0\x1e\xca\x3f\x5e\x06\x9e\x3f"
|
|
val2 := &redis.VectorFP32{
|
|
Val: []byte(fp32),
|
|
}
|
|
ok, err = client.VAdd(ctx, vecName, "k2", val2).Result()
|
|
expectNil(err)
|
|
expectTrue(ok)
|
|
|
|
dim, err := client.VDim(ctx, vecName).Result()
|
|
expectNil(err)
|
|
expectEqual(dim, 4)
|
|
|
|
count, err := client.VCard(ctx, vecName).Result()
|
|
expectNil(err)
|
|
expectEqual(count, 2)
|
|
|
|
ok, err = client.VRem(ctx, vecName, "k1").Result()
|
|
expectNil(err)
|
|
expectTrue(ok)
|
|
|
|
count, err = client.VCard(ctx, vecName).Result()
|
|
expectNil(err)
|
|
expectEqual(count, 1)
|
|
})
|
|
|
|
It("basic similarity", func() {
|
|
SkipBeforeRedisVersion(8.0, "Redis 8.0 introduces support for VectorSet")
|
|
vecName := "basic_similarity"
|
|
|
|
ok, err := client.VAdd(ctx, vecName, "k1", &redis.VectorValues{
|
|
Val: []float64{1, 0, 0, 0},
|
|
}).Result()
|
|
expectNil(err)
|
|
expectTrue(ok)
|
|
ok, err = client.VAdd(ctx, vecName, "k2", &redis.VectorValues{
|
|
Val: []float64{0.99, 0.01, 0, 0},
|
|
}).Result()
|
|
expectNil(err)
|
|
expectTrue(ok)
|
|
ok, err = client.VAdd(ctx, vecName, "k3", &redis.VectorValues{
|
|
Val: []float64{0.1, 1, -1, 0.5},
|
|
}).Result()
|
|
expectNil(err)
|
|
expectTrue(ok)
|
|
|
|
sim, err := client.VSimWithScores(ctx, vecName, &redis.VectorValues{
|
|
Val: []float64{1, 0, 0, 0},
|
|
}).Result()
|
|
expectNil(err)
|
|
expectEqual(len(sim), 3)
|
|
simMap := make(map[string]float64)
|
|
for _, vi := range sim {
|
|
simMap[vi.Name] = vi.Score
|
|
}
|
|
expectTrue(simMap["k1"] > 0.99)
|
|
expectTrue(simMap["k2"] > 0.99)
|
|
expectTrue(simMap["k3"] < 0.8)
|
|
})
|
|
|
|
It("dimension operation", func() {
|
|
SkipBeforeRedisVersion(8.0, "Redis 8.0 introduces support for VectorSet")
|
|
vecName := "dimension_op"
|
|
originalDim := 100
|
|
reducedDim := 50
|
|
|
|
v1 := generateRandomVector(originalDim)
|
|
ok, err := client.VAddWithArgs(ctx, vecName, "k1", &v1, &redis.VAddArgs{
|
|
Reduce: int64(reducedDim),
|
|
}).Result()
|
|
expectNil(err)
|
|
expectTrue(ok)
|
|
|
|
info, err := client.VInfo(ctx, vecName).Result()
|
|
expectNil(err)
|
|
dim := info["vector-dim"].(int64)
|
|
oriDim := info["projection-input-dim"].(int64)
|
|
expectEqual(dim, reducedDim)
|
|
expectEqual(oriDim, originalDim)
|
|
|
|
wrongDim := 80
|
|
wrongV := generateRandomVector(wrongDim)
|
|
_, err = client.VAddWithArgs(ctx, vecName, "kw", &wrongV, &redis.VAddArgs{
|
|
Reduce: int64(reducedDim),
|
|
}).Result()
|
|
expectTrue(err != nil)
|
|
|
|
v2 := generateRandomVector(originalDim)
|
|
ok, err = client.VAddWithArgs(ctx, vecName, "k2", &v2, &redis.VAddArgs{
|
|
Reduce: int64(reducedDim),
|
|
}).Result()
|
|
expectNil(err)
|
|
expectTrue(ok)
|
|
})
|
|
|
|
It("remove", func() {
|
|
SkipBeforeRedisVersion(8.0, "Redis 8.0 introduces support for VectorSet")
|
|
vecName := "remove"
|
|
v1 := generateRandomVector(5)
|
|
ok, err := client.VAdd(ctx, vecName, "k1", &v1).Result()
|
|
expectNil(err)
|
|
expectTrue(ok)
|
|
|
|
exist, err := client.Exists(ctx, vecName).Result()
|
|
expectNil(err)
|
|
expectEqual(exist, 1)
|
|
|
|
ok, err = client.VRem(ctx, vecName, "k1").Result()
|
|
expectNil(err)
|
|
expectTrue(ok)
|
|
|
|
exist, err = client.Exists(ctx, vecName).Result()
|
|
expectNil(err)
|
|
expectEqual(exist, 0)
|
|
})
|
|
|
|
It("all operations", func() {
|
|
SkipBeforeRedisVersion(8.0, "Redis 8.0 introduces support for VectorSet")
|
|
vecName := "commands"
|
|
vals := []struct {
|
|
name string
|
|
v redis.VectorValues
|
|
attr string
|
|
}{
|
|
{
|
|
name: "k0",
|
|
v: redis.VectorValues{Val: []float64{1, 0, 0, 0}},
|
|
attr: `{"age": 25, "name": "Alice", "active": true, "scores": [85, 90, 95], "city": "New York"}`,
|
|
},
|
|
{
|
|
name: "k1",
|
|
v: redis.VectorValues{Val: []float64{0, 1, 0, 0}},
|
|
attr: `{"age": 30, "name": "Bob", "active": false, "scores": [70, 75, 80], "city": "Boston"}`,
|
|
},
|
|
{
|
|
name: "k2",
|
|
v: redis.VectorValues{Val: []float64{0, 0, 1, 0}},
|
|
attr: `{"age": 35, "name": "Charlie", "scores": [60, 65, 70], "city": "Seattle"}`,
|
|
},
|
|
{
|
|
name: "k3",
|
|
v: redis.VectorValues{Val: []float64{0, 0, 0, 1}},
|
|
},
|
|
{
|
|
name: "k4",
|
|
v: redis.VectorValues{Val: []float64{0.5, 0.5, 0, 0}},
|
|
attr: `invalid json`,
|
|
},
|
|
}
|
|
|
|
// If the key doesn't exist, return null error
|
|
_, err := client.VRandMember(ctx, vecName).Result()
|
|
expectEqual(err.Error(), proto.Nil.Error())
|
|
|
|
// If the key doesn't exist, return an empty array
|
|
res, err := client.VRandMemberCount(ctx, vecName, 3).Result()
|
|
expectNil(err)
|
|
expectEqual(len(res), 0)
|
|
|
|
for _, v := range vals {
|
|
ok, err := client.VAdd(ctx, vecName, v.name, &v.v).Result()
|
|
expectNil(err)
|
|
expectTrue(ok)
|
|
if len(v.attr) > 0 {
|
|
ok, err = client.VSetAttr(ctx, vecName, v.name, v.attr).Result()
|
|
expectNil(err)
|
|
expectTrue(ok)
|
|
}
|
|
}
|
|
|
|
// VGetAttr
|
|
attr, err := client.VGetAttr(ctx, vecName, vals[1].name).Result()
|
|
expectNil(err)
|
|
expectEqual(attr, vals[1].attr)
|
|
|
|
// VRandMember
|
|
_, err = client.VRandMember(ctx, vecName).Result()
|
|
expectNil(err)
|
|
|
|
res, err = client.VRandMemberCount(ctx, vecName, 3).Result()
|
|
expectNil(err)
|
|
expectEqual(len(res), 3)
|
|
|
|
res, err = client.VRandMemberCount(ctx, vecName, 10).Result()
|
|
expectNil(err)
|
|
expectEqual(len(res), len(vals))
|
|
|
|
// test equality
|
|
sim, err := client.VSimWithArgs(ctx, vecName, &vals[0].v, &redis.VSimArgs{
|
|
Filter: `.age == 25`,
|
|
}).Result()
|
|
expectNil(err)
|
|
expectEqual(len(sim), 1)
|
|
expectEqual(sim[0], vals[0].name)
|
|
|
|
// test greater than
|
|
sim, err = client.VSimWithArgs(ctx, vecName, &vals[0].v, &redis.VSimArgs{
|
|
Filter: `.age > 25`,
|
|
}).Result()
|
|
expectNil(err)
|
|
expectEqual(len(sim), 2)
|
|
|
|
// test less than or equal
|
|
sim, err = client.VSimWithArgs(ctx, vecName, &vals[0].v, &redis.VSimArgs{
|
|
Filter: `.age <= 30`,
|
|
}).Result()
|
|
expectNil(err)
|
|
expectEqual(len(sim), 2)
|
|
|
|
// test string equality
|
|
sim, err = client.VSimWithArgs(ctx, vecName, &vals[0].v, &redis.VSimArgs{
|
|
Filter: `.name == "Alice"`,
|
|
}).Result()
|
|
expectNil(err)
|
|
expectEqual(len(sim), 1)
|
|
expectEqual(sim[0], vals[0].name)
|
|
|
|
// test string inequality
|
|
sim, err = client.VSimWithArgs(ctx, vecName, &vals[0].v, &redis.VSimArgs{
|
|
Filter: `.name != "Alice"`,
|
|
}).Result()
|
|
expectNil(err)
|
|
expectEqual(len(sim), 2)
|
|
|
|
// test bool
|
|
sim, err = client.VSimWithArgs(ctx, vecName, &vals[0].v, &redis.VSimArgs{
|
|
Filter: `.active`,
|
|
}).Result()
|
|
expectNil(err)
|
|
expectEqual(len(sim), 1)
|
|
expectEqual(sim[0], vals[0].name)
|
|
|
|
// test logical add
|
|
sim, err = client.VSimWithArgs(ctx, vecName, &vals[0].v, &redis.VSimArgs{
|
|
Filter: `.age > 20 and .age < 30`,
|
|
}).Result()
|
|
expectNil(err)
|
|
expectEqual(len(sim), 1)
|
|
expectEqual(sim[0], vals[0].name)
|
|
|
|
// test logical or
|
|
sim, err = client.VSimWithArgs(ctx, vecName, &vals[0].v, &redis.VSimArgs{
|
|
Filter: `.age < 30 or .age > 35`,
|
|
}).Result()
|
|
expectNil(err)
|
|
expectEqual(len(sim), 1)
|
|
expectEqual(sim[0], vals[0].name)
|
|
})
|
|
})
|
|
}
|
|
})
|