Skip to content

feat: support vectorset #3375

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 55 additions & 0 deletions command.go
Original file line number Diff line number Diff line change
Expand Up @@ -5615,3 +5615,58 @@ func (cmd *MonitorCmd) Stop() {
defer cmd.mu.Unlock()
cmd.status = monitorStatusStop
}

type MapStringFloatCmd struct {
baseCmd

val map[string]float64
}

var _ Cmder = (*MapStringFloatCmd)(nil)

func NewMapStringFloatCmd(ctx context.Context, args ...interface{}) *MapStringFloatCmd {
return &MapStringFloatCmd{
baseCmd: baseCmd{
ctx: ctx,
args: args,
},
}
}

func (cmd *MapStringFloatCmd) SetVal(val map[string]float64) {
cmd.val = val
}

func (cmd *MapStringFloatCmd) Val() map[string]float64 {
return cmd.val
}

func (cmd *MapStringFloatCmd) Result() (map[string]float64, error) {
return cmd.val, cmd.err
}

func (cmd *MapStringFloatCmd) String() string {
return cmdString(cmd, cmd.val)
}

func (cmd *MapStringFloatCmd) readReply(rd *proto.Reader) error {
n, err := rd.ReadMapLen()
if err != nil {
return err
}

cmd.val = make(map[string]float64, n)
for i := 0; i < n; i++ {
key, err := rd.ReadString()
if err != nil {
return err
}

nn, err := rd.ReadFloat()
if err != nil {
return err
}
cmd.val[key] = nn
}
return nil
}
1 change: 1 addition & 0 deletions commands.go
Original file line number Diff line number Diff line change
Expand Up @@ -234,6 +234,7 @@ type Cmdable interface {
StreamCmdable
TimeseriesCmdable
JSONCmdable
VectorSetCmdable
}

type StatefulCmdable interface {
Expand Down
284 changes: 284 additions & 0 deletions vectorset_commands.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,284 @@
package redis

import (
"context"
"strconv"
)

type VectorSetCmdable interface {
VAdd(ctx context.Context, key, element string, val Vector) *BoolCmd
VAddArgs(ctx context.Context, key, element string, val Vector, addArgs VAddArgs) *BoolCmd
VCard(ctx context.Context, key string) *IntCmd
VDim(ctx context.Context, key string) *IntCmd
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) *MapStringFloatCmd
VRandMember(ctx context.Context, key string) *StringCmd
VRandMemberCount(ctx context.Context, key string, count int) *StringSliceCmd
VRem(ctx context.Context, key, element string) *BoolCmd
VSetAttr(ctx context.Context, key, element, attr string) *BoolCmd
VSim(ctx context.Context, key string, val Vector) *StringSliceCmd
VSimWithScores(ctx context.Context, key string, val Vector) *MapStringFloatCmd
VSimArgs(ctx context.Context, key string, val Vector, args VSimArgs) *StringSliceCmd
VSimArgsWithScores(ctx context.Context, key string, val Vector, args VSimArgs) *MapStringFloatCmd
}

type Vector interface {
Value() []any
}

const (
vectorFormatFP32 string = "FP32"
vectorFormatValues string = "Values"
)

type VectorFP32 struct {
Val []byte
}

func (v *VectorFP32) Value() []any {
return []any{vectorFormatFP32, v.Val}
}

var _ Vector = &VectorFP32{}

type VectorValues struct {
Val []float64
}

func (v *VectorValues) Value() []any {
res := make([]any, 2+len(v.Val))
res[0] = vectorFormatValues
res[1] = len(v.Val)
for i, v := range v.Val {
res[2+i] = v
}
return res
}

var _ Vector = &VectorValues{}

type VectorRef struct {
Name string // the name of the referent vector
}

func (v *VectorRef) Value() []any {
return []any{"ele", v.Name}
}

var _ Vector = &VectorRef{}

// `VADD key (FP32 | VALUES num) vector element`
func (c cmdable) VAdd(ctx context.Context, key, element string, val Vector) *BoolCmd {
return c.VAddArgs(ctx, key, element, val, VAddArgs{})
}

type VAddArgs struct {
// the REDUCE option must be passed immediately after the key
Reduce int64
Cas bool

// The NoQuant, Q8 and Bin options are mutually exclusive.
NoQuant bool
Q8 bool
Bin bool

EF int64
SetAttr string
M int64
}

func (v VAddArgs) reduce() int64 {
return v.Reduce
}

func (v VAddArgs) appendArgs(args []any) []any {
if v.Cas {
args = append(args, "cas")
}

if v.NoQuant {
args = append(args, "noquant")
} else if v.Q8 {
args = append(args, "q8")
} else if v.Bin {
args = append(args, "bin")
}

if v.EF > 0 {
args = append(args, "ef", strconv.FormatInt(v.EF, 10))
}
if len(v.SetAttr) > 0 {
args = append(args, "setattr", v.SetAttr)
}
if v.M > 0 {
args = append(args, "m", strconv.FormatInt(v.M, 10))
}
return args
}

// `VADD key [REDUCE dim] (FP32 | VALUES num) vector element [CAS] [NOQUANT | Q8 | BIN] [EF build-exploration-factor] [SETATTR attributes] [M numlinks]`
func (c cmdable) VAddArgs(ctx context.Context, key, element string, val Vector, addArgs VAddArgs) *BoolCmd {
args := []any{"vadd", key}
if addArgs.reduce() > 0 {
args = append(args, "reduce", addArgs.reduce())
}
args = append(args, val.Value()...)
args = append(args, element)
args = addArgs.appendArgs(args)
cmd := NewBoolCmd(ctx, args...)
_ = c(ctx, cmd)
return cmd
}

// `VCARD key`
func (c cmdable) VCard(ctx context.Context, key string) *IntCmd {
cmd := NewIntCmd(ctx, "vcard", key)
_ = c(ctx, cmd)
return cmd
}

// `VDIM key`
func (c cmdable) VDim(ctx context.Context, key string) *IntCmd {
cmd := NewIntCmd(ctx, "vdim", key)
_ = c(ctx, cmd)
return cmd
}

// `VEMB key element [RAW]`
func (c cmdable) VEmb(ctx context.Context, key, element string, raw bool) *SliceCmd {
args := []any{"vemb", key, element}
if raw {
args = append(args, "raw")
}
cmd := NewSliceCmd(ctx, args...)
_ = c(ctx, cmd)
return cmd
}

// `VGETATTR key element`
func (c cmdable) VGetAttr(ctx context.Context, key, element string) *StringCmd {
cmd := NewStringCmd(ctx, "vgetattr", key, element)
_ = c(ctx, cmd)
return cmd
}

// `VINFO key`
func (c cmdable) VInfo(ctx context.Context, key string) *MapStringInterfaceCmd {
cmd := NewMapStringInterfaceCmd(ctx, "vinfo", key)
_ = c(ctx, cmd)
return cmd
}

// `VLINKS key element`
func (c cmdable) VLinks(ctx context.Context, key, element string) *StringSliceCmd {
cmd := NewStringSliceCmd(ctx, "vlinks", key, element)
_ = c(ctx, cmd)
return cmd
}

// `VLINKS key element WITHSCORES`
func (c cmdable) VLinksWithScores(ctx context.Context, key, element string) *MapStringFloatCmd {
cmd := NewMapStringFloatCmd(ctx, "vlinks", key, element, "withscores")
_ = c(ctx, cmd)
return cmd
}

// `VRANDMEMBER key`
func (c cmdable) VRandMember(ctx context.Context, key string) *StringCmd {
cmd := NewStringCmd(ctx, "vrandmember", key)
_ = c(ctx, cmd)
return cmd
}

// `VRANDMEMBER key [count]`
func (c cmdable) VRandMemberCount(ctx context.Context, key string, count int) *StringSliceCmd {
cmd := NewStringSliceCmd(ctx, "vrandmember", key, count)
_ = c(ctx, cmd)
return cmd
}

// `VREM key element`
func (c cmdable) VRem(ctx context.Context, key, element string) *BoolCmd {
cmd := NewBoolCmd(ctx, "vrem", key, element)
_ = c(ctx, cmd)
return cmd
}

// `VSETATTR key element "{ JSON obj }"`
func (c cmdable) VSetAttr(ctx context.Context, key, element, attr string) *BoolCmd {
cmd := NewBoolCmd(ctx, "vsetattr", key, element, attr)
_ = c(ctx, cmd)
return cmd
}

// `VSIM key (ELE | FP32 | VALUES num) (vector | element)`
func (c cmdable) VSim(ctx context.Context, key string, val Vector) *StringSliceCmd {
return c.VSimArgs(ctx, key, val, VSimArgs{})
}

// `VSIM key (ELE | FP32 | VALUES num) (vector | element) WITHSCORES`
func (c cmdable) VSimWithScores(ctx context.Context, key string, val Vector) *MapStringFloatCmd {
return c.VSimArgsWithScores(ctx, key, val, VSimArgs{})
}

type VSimArgs struct {
Count int64
EF int64
Filter string
FilterEF int64
Truth bool
NoThread bool
// The `VSim` command in Redis has the option, by the doc in Redis.io don't have.
// Epsilon float64
}

func (v VSimArgs) appendArgs(args []any) []any {
if v.Count > 0 {
args = append(args, "count", v.Count)
}
if v.EF > 0 {
args = append(args, "ef", v.EF)
}
if len(v.Filter) > 0 {
args = append(args, "filter", v.Filter)
}
if v.FilterEF > 0 {
args = append(args, "filter-ef", v.FilterEF)
}
if v.Truth {
args = append(args, "truth")
}
if v.NoThread {
args = append(args, "nothread")
}
// if v.Epsilon > 0 {
// args = append(args, "Epsilon", v.Epsilon)
// }
return args
}

// `VSIM key (ELE | FP32 | VALUES num) (vector | element) [COUNT num]
// [EF search-exploration-factor] [FILTER expression] [FILTER-EF max-filtering-effort] [TRUTH] [NOTHREAD]`
func (c cmdable) VSimArgs(ctx context.Context, key string, val Vector, simArgs VSimArgs) *StringSliceCmd {
args := []any{"vsim", key}
args = append(args, val.Value()...)
args = simArgs.appendArgs(args)
cmd := NewStringSliceCmd(ctx, args...)
_ = c(ctx, cmd)
return cmd
}

// `VSIM key (ELE | FP32 | VALUES num) (vector | element) [WITHSCORES] [COUNT num]
// [EF search-exploration-factor] [FILTER expression] [FILTER-EF max-filtering-effort] [TRUTH] [NOTHREAD]`
func (c cmdable) VSimArgsWithScores(ctx context.Context, key string, val Vector, simArgs VSimArgs) *MapStringFloatCmd {
args := []any{"vsim", key}
args = append(args, val.Value()...)
args = append(args, "withscores")
args = simArgs.appendArgs(args)
cmd := NewMapStringFloatCmd(ctx, args...)
_ = c(ctx, cmd)
return cmd
}
Loading
Loading