diff --git a/command.go b/command.go index 6dec765a..7cea6784 100644 --- a/command.go +++ b/command.go @@ -1763,6 +1763,87 @@ func (cmd *StringSliceCmd) Clone() Cmder { //------------------------------------------------------------------------------ +// StringSliceSliceCmd returns a slice of string slices ([][]string). +// This is used for commands like VLINKS that return an array of arrays. +type StringSliceSliceCmd struct { + baseCmd + + val [][]string +} + +var _ Cmder = (*StringSliceSliceCmd)(nil) + +func NewStringSliceSliceCmd(ctx context.Context, args ...any) *StringSliceSliceCmd { + return &StringSliceSliceCmd{ + baseCmd: baseCmd{ + ctx: ctx, + args: args, + }, + } +} + +func (cmd *StringSliceSliceCmd) SetVal(val [][]string) { + cmd.val = val +} + +func (cmd *StringSliceSliceCmd) Val() [][]string { + return cmd.val +} + +func (cmd *StringSliceSliceCmd) Result() ([][]string, error) { + return cmd.val, cmd.err +} + +func (cmd *StringSliceSliceCmd) String() string { + return cmdString(cmd, cmd.val) +} + +func (cmd *StringSliceSliceCmd) readReply(rd *proto.Reader) error { + n, err := rd.ReadArrayLen() + if err != nil { + return err + } + cmd.val = make([][]string, n) + for i := range n { + // Read inner array + innerN, err := rd.ReadArrayLen() + if err != nil { + return err + } + cmd.val[i] = make([]string, innerN) + for j := range innerN { + switch s, err := rd.ReadString(); { + case err == Nil: + cmd.val[i][j] = "" + case err != nil: + return err + default: + cmd.val[i][j] = s + } + } + } + return nil +} + +func (cmd *StringSliceSliceCmd) Clone() Cmder { + var val [][]string + if cmd.val != nil { + val = make([][]string, len(cmd.val)) + for i, slice := range cmd.val { + if slice != nil { + val[i] = make([]string, len(slice)) + copy(val[i], slice) + } + } + } + return &StringSliceSliceCmd{ + baseCmd: cmd.cloneBaseCmd(), + val: val, + } +} + +//------------------------------------------------------------------------------ + type KeyValue struct { Key string Value string @@ -7942,6 +8023,120 @@ func (cmd *VectorScoreSliceCmd) Clone() Cmder { } } +// VectorScoreSliceSliceCmd is used for VLINKS WITHSCORES which returns an array of arrays. +// In RESP3, each inner array contains maps of element -> score. +type VectorScoreSliceSliceCmd struct { + baseCmd + + val [][]VectorScore +} + +var _ Cmder = (*VectorScoreSliceSliceCmd)(nil) + +func NewVectorScoreSliceSliceCmd(ctx context.Context, args ...any) *VectorScoreSliceSliceCmd { + return &VectorScoreSliceSliceCmd{ + baseCmd: baseCmd{ + ctx: ctx, + args: args, + }, + } +} + +func (cmd *VectorScoreSliceSliceCmd) SetVal(val [][]VectorScore) { + cmd.val = val +} + +func (cmd *VectorScoreSliceSliceCmd) Val() [][]VectorScore { + return cmd.val +} + +func (cmd *VectorScoreSliceSliceCmd) Result() ([][]VectorScore, error) { + return cmd.val, cmd.err +} + +func (cmd *VectorScoreSliceSliceCmd) String() string { + return cmdString(cmd, cmd.val) +} + +func (cmd *VectorScoreSliceSliceCmd) readReply(rd *proto.Reader) error { + n, err := rd.ReadArrayLen() + if err != nil { + return err + } + + cmd.val = make([][]VectorScore, n) + for i := range n { + // Each level can be either a map (RESP3) or an array (RESP2) + levelTyp, err := rd.PeekReplyType() + if err != nil { + return err + } + + if levelTyp == proto.RespMap { + // RESP3 format: each level is a map {element: score, element: score, ...} + mapLen, err := rd.ReadMapLen() + if err != nil { + return err + } + + cmd.val[i] = make([]VectorScore, mapLen) + for j := range mapLen { + name, err := rd.ReadString() + if err != nil { + return err + } + score, err := rd.ReadFloat() + if err != nil { + return err + } + cmd.val[i][j] = VectorScore{Name: name, Score: score} + } + } else { + // RESP2 format: each level is an array of [element, score, element, score, ...] pairs + innerLen, err := rd.ReadArrayLen() + if err != nil { + return err + } + + if innerLen%2 != 0 { + return fmt.Errorf("redis: got %d elements in the VLINKS array, wanted a multiple of 2", innerLen) + } + + cmd.val[i] = make([]VectorScore, innerLen/2) + for j := 0; j < innerLen; j += 2 { + name, err := rd.ReadString() + if err != nil { + return err + } + score, err := rd.ReadFloat() + if err != nil { + return err + } + cmd.val[i][j/2] = VectorScore{Name: name, Score: score} + } + } + } + + return nil +} + +func (cmd *VectorScoreSliceSliceCmd) Clone() Cmder { + var val [][]VectorScore + if cmd.val != nil { + val = make([][]VectorScore, len(cmd.val)) + for i, slice := range cmd.val { + if slice != nil { + val[i] = make([]VectorScore, len(slice)) + copy(val[i], slice) + } + } + } + return &VectorScoreSliceSliceCmd{ + baseCmd: cmd.cloneBaseCmd(), + val: val, + } +} + func readVectorAttribStringOrNil(rd *proto.Reader) (*string, error) { v, err := rd.ReadReply() if err != nil { diff --git a/doctests/vec_set_test.go b/doctests/vec_set_test.go index 1788ca2d..ebd0ca61 100644 --- a/doctests/vec_set_test.go +++ b/doctests/vec_set_test.go @@ -385,6 +385,24 @@ func ExampleClient_vectorset() { fmt.Println(res34[0].Name, res34[0].Score, *res34[0].Attribs) // >>> pt:A 1 {"price":18.99,"size":"large"} // STEP_END + // STEP_START vlinks + res35, err := rdb.VLinks(ctx, "points", "pt:A").Result() + if err != nil { + panic(err) + } + + // Print only the last layer (deterministic) + fmt.Println(res35[len(res35)-1]) // [pt:B pt:C pt:D pt:E] + + res36, err := rdb.VLinksWithScores(ctx, "points", "pt:A").Result() + if err != nil { + panic(err) + } + + // Print only the last layer (deterministic) + fmt.Println(res36[len(res36)-1]) // [{pt:B 0} {pt:C 0.5} {pt:D 0.5} {pt:E 0.8535534143447876}] + // STEP_END + // Output: // true // true @@ -420,6 +438,8 @@ func ExampleClient_vectorset() { // false // pt:A {"price":18.99,"size":"large"} // pt:A 1 {"price":18.99,"size":"large"} + // [pt:B pt:C pt:D pt:E] + // [{pt:B 0} {pt:C 0.5} {pt:D 0.5} {pt:E 0.8535534143447876}] } func ExampleClient_vectorset_quantization() { diff --git a/vectorset_commands.go b/vectorset_commands.go index 5516da67..a88bb1cd 100644 --- a/vectorset_commands.go +++ b/vectorset_commands.go @@ -15,8 +15,8 @@ type VectorSetCmdable interface { VEmb(ctx context.Context, key, element string, raw bool) *SliceCmd VGetAttr(ctx context.Context, key, element string) *StringCmd VInfo(ctx context.Context, key string) *MapStringInterfaceCmd - VLinks(ctx context.Context, key, element string) *StringSliceCmd - VLinksWithScores(ctx context.Context, key, element string) *VectorScoreSliceCmd + VLinks(ctx context.Context, key, element string) *StringSliceSliceCmd + VLinksWithScores(ctx context.Context, key, element string) *VectorScoreSliceSliceCmd VRandMember(ctx context.Context, key string) *StringCmd VRandMemberCount(ctx context.Context, key string, count int) *StringSliceCmd VRem(ctx context.Context, key, element string) *BoolCmd @@ -272,16 +272,16 @@ func (c cmdable) VInfo(ctx context.Context, key string) *MapStringInterfaceCmd { // `VLINKS key element` // note: the API is experimental and may be subject to change. -func (c cmdable) VLinks(ctx context.Context, key, element string) *StringSliceCmd { - cmd := NewStringSliceCmd(ctx, "vlinks", key, element) +func (c cmdable) VLinks(ctx context.Context, key, element string) *StringSliceSliceCmd { + cmd := NewStringSliceSliceCmd(ctx, "vlinks", key, element) _ = c(ctx, cmd) return cmd } // `VLINKS key element WITHSCORES` // note: the API is experimental and may be subject to change. -func (c cmdable) VLinksWithScores(ctx context.Context, key, element string) *VectorScoreSliceCmd { - cmd := NewVectorInfoSliceCmd(ctx, "vlinks", key, element, "withscores") +func (c cmdable) VLinksWithScores(ctx context.Context, key, element string) *VectorScoreSliceSliceCmd { + cmd := NewVectorScoreSliceSliceCmd(ctx, "vlinks", key, element, "withscores") _ = c(ctx, cmd) return cmd } diff --git a/vectorset_commands_integration_test.go b/vectorset_commands_integration_test.go index d4e4f302..247d6762 100644 --- a/vectorset_commands_integration_test.go +++ b/vectorset_commands_integration_test.go @@ -394,6 +394,43 @@ var _ = Describe("Redis VectorSet commands", Label("vectorset"), func() { expectEqual(simScoresAttribs[0].Score, float64(1)) expectEqual(simScoresAttribs[0].Attribs, (*string)(nil)) } + + // test vlinks and vlinks with scores + links, err := client.VLinks(ctx, vecName, "k1").Result() + expectNil(err) + expectTrue(len(links) > 0) + + // Last layer should contain all other elements (order may vary) + lastLayer := links[len(links)-1] + expectEqual(len(lastLayer), 4) + + linksWithScores, err := client.VLinksWithScores(ctx, vecName, "k1").Result() + expectNil(err) + expectTrue(len(linksWithScores) > 0) + + // Last layer should contain all other elements with scores + lastLayerScores := linksWithScores[len(linksWithScores)-1] + expectEqual(len(lastLayerScores), 4) + }) + + It("vlinks with single element", func() { + SkipBeforeRedisVersion(8.0, "Redis 8.0 introduces support for VectorSet") + vecName := "vlinks_single" + + // Add only one vector + ok, err := client.VAdd(ctx, vecName, "only", &redis.VectorValues{ + Val: []float64{1, 1}, + }).Result() + expectNil(err) + expectTrue(ok) + + links, err := client.VLinks(ctx, vecName, "only").Result() + expectNil(err) + expectTrue(len(links) > 0) + + linksWithScores, err := client.VLinksWithScores(ctx, vecName, "only").Result() + expectNil(err) + expectTrue(len(linksWithScores) > 0) }) }) } diff --git a/vectorset_commands_test.go b/vectorset_commands_test.go index 2131e192..441ef160 100644 --- a/vectorset_commands_test.go +++ b/vectorset_commands_test.go @@ -171,7 +171,7 @@ func TestVLinks(t *testing.T) { m := &mockCmdable{} c := m.asCmdable() c.VLinks(context.Background(), "k", "e") - cmd := m.lastCmd.(*StringSliceCmd) + cmd := m.lastCmd.(*StringSliceSliceCmd) if cmd.args[0] != "vlinks" || cmd.args[1] != "k" || cmd.args[2] != "e" { t.Errorf("unexpected args: %v", cmd.args) } @@ -181,7 +181,7 @@ func TestVLinksWithScores(t *testing.T) { m := &mockCmdable{} c := m.asCmdable() c.VLinksWithScores(context.Background(), "k", "e") - cmd := m.lastCmd.(*VectorScoreSliceCmd) + cmd := m.lastCmd.(*VectorScoreSliceSliceCmd) if cmd.args[0] != "vlinks" || cmd.args[1] != "k" || cmd.args[2] != "e" || cmd.args[3] != "withscores" { t.Errorf("unexpected args: %v", cmd.args) } @@ -826,3 +826,122 @@ func TestVectorScoreSliceCmdReadReply(t *testing.T) { }) } } + +// TestStringSliceSliceCmdReadReply covers the readReply parser used by VLINKS, +// which returns an array of arrays of element names (one inner array per HNSW layer). +func TestStringSliceSliceCmdReadReply(t *testing.T) { + tests := []struct { + name string + respData []byte + want [][]string + wantErr bool + }{ + { + name: "two layers with elements", + respData: []byte("*2\r\n*2\r\n$2\r\ne1\r\n$2\r\ne2\r\n*1\r\n$2\r\ne3\r\n"), + want: [][]string{{"e1", "e2"}, {"e3"}}, + }, + { + name: "empty outer array", + respData: []byte("*0\r\n"), + want: [][]string{}, + }, + { + name: "outer array with empty inner", + respData: []byte("*1\r\n*0\r\n"), + want: [][]string{{}}, + }, + { + name: "nil element becomes empty string", + respData: []byte("*1\r\n*2\r\n$2\r\ne1\r\n$-1\r\n"), + want: [][]string{{"e1", ""}}, + }, + { + name: "malformed outer length", + respData: []byte("+notarray\r\n"), + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rd := proto.NewReader(newMockConn(tt.respData)) + cmd := NewStringSliceSliceCmd(context.Background(), "vlinks", "key", "e") + err := cmd.readReply(rd) + if (err != nil) != tt.wantErr { + t.Fatalf("readReply() error = %v, wantErr %v", err, tt.wantErr) + } + if !tt.wantErr && !reflect.DeepEqual(cmd.Val(), tt.want) { + t.Errorf("Val() = %v, want %v", cmd.Val(), tt.want) + } + }) + } +} + +// TestVectorScoreSliceSliceCmdReadReply covers the readReply parser used by +// VLINKS WITHSCORES, which returns an array of layers; each layer is a map +// (RESP3) or flat element/score array (RESP2). +func TestVectorScoreSliceSliceCmdReadReply(t *testing.T) { + tests := []struct { + name string + respData []byte + want [][]VectorScore + wantErr bool + }{ + { + name: "RESP3 two layers, maps", + respData: []byte("*2\r\n%2\r\n+a1\r\n,0.9\r\n+a2\r\n,0.8\r\n%1\r\n+b1\r\n,0.5\r\n"), + want: [][]VectorScore{ + {{Name: "a1", Score: 0.9}, {Name: "a2", Score: 0.8}}, + {{Name: "b1", Score: 0.5}}, + }, + }, + { + name: "RESP2 two layers, flat arrays", + respData: []byte("*2\r\n*4\r\n$2\r\na1\r\n$3\r\n0.9\r\n$2\r\na2\r\n$3\r\n0.8\r\n*2\r\n$2\r\nb1\r\n$3\r\n0.5\r\n"), + want: [][]VectorScore{ + {{Name: "a1", Score: 0.9}, {Name: "a2", Score: 0.8}}, + {{Name: "b1", Score: 0.5}}, + }, + }, + { + name: "RESP3 empty layer (empty map)", + respData: []byte("*1\r\n%0\r\n"), + want: [][]VectorScore{{}}, + }, + { + name: "RESP2 empty layer (empty array)", + respData: []byte("*1\r\n*0\r\n"), + want: [][]VectorScore{{}}, + }, + { + name: "empty outer", + respData: []byte("*0\r\n"), + want: [][]VectorScore{}, + }, + { + name: "RESP2 odd inner array returns error", + respData: []byte("*1\r\n*3\r\n$2\r\na1\r\n$3\r\n0.9\r\n$2\r\na2\r\n"), + wantErr: true, + }, + { + name: "malformed outer length", + respData: []byte("+notarray\r\n"), + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rd := proto.NewReader(newMockConn(tt.respData)) + cmd := NewVectorScoreSliceSliceCmd(context.Background(), "vlinks", "key", "e", "withscores") + err := cmd.readReply(rd) + if (err != nil) != tt.wantErr { + t.Fatalf("readReply() error = %v, wantErr %v", err, tt.wantErr) + } + if !tt.wantErr && !reflect.DeepEqual(cmd.Val(), tt.want) { + t.Errorf("Val() = %v, want %v", cmd.Val(), tt.want) + } + }) + } +}