Skip to content

Commit 3c85d09

Browse files
ofekshenawaofekshenawa
authored
feat(search): Add VAMANA vector type to RediSearch (#3449)
* Add VAMANA vector type to redisearch * Change to svs-vamana vector type && remove panics from search module * fix tests * fix tests * fix tests
1 parent 9177a76 commit 3c85d09

File tree

2 files changed

+753
-28
lines changed

2 files changed

+753
-28
lines changed

‎search_commands.go‎

Lines changed: 105 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -80,8 +80,9 @@ type FieldSchema struct {
8080
}
8181

8282
type FTVectorArgs struct {
83-
FlatOptions *FTFlatOptions
84-
HNSWOptions *FTHNSWOptions
83+
FlatOptions *FTFlatOptions
84+
HNSWOptions *FTHNSWOptions
85+
VamanaOptions *FTVamanaOptions
8586
}
8687

8788
type FTFlatOptions struct {
@@ -103,6 +104,19 @@ type FTHNSWOptions struct {
103104
Epsilon float64
104105
}
105106

107+
type FTVamanaOptions struct {
108+
Type string
109+
Dim int
110+
DistanceMetric string
111+
Compression string
112+
ConstructionWindowSize int
113+
GraphMaxDegree int
114+
SearchWindowSize int
115+
Epsilon float64
116+
TrainingThreshold int
117+
ReduceDim int
118+
}
119+
106120
type FTDropIndexOptions struct {
107121
DeleteDocs bool
108122
}
@@ -499,7 +513,7 @@ func (c cmdable) FTAggregate(ctx context.Context, index string, query string) *M
499513
return cmd
500514
}
501515

502-
func FTAggregateQuery(query string, options *FTAggregateOptions) AggregateQuery {
516+
func FTAggregateQuery(query string, options *FTAggregateOptions) (AggregateQuery, error) {
503517
queryArgs := []interface{}{query}
504518
if options != nil {
505519
if options.Verbatim {
@@ -515,7 +529,7 @@ func FTAggregateQuery(query string, options *FTAggregateOptions) AggregateQuery
515529
}
516530

517531
if options.LoadAll && options.Load != nil {
518-
panic("FT.AGGREGATE: LOADALL and LOAD are mutually exclusive")
532+
return nil, fmt.Errorf("FT.AGGREGATE: LOADALL and LOAD are mutually exclusive")
519533
}
520534
if options.LoadAll {
521535
queryArgs = append(queryArgs, "LOAD", "*")
@@ -571,7 +585,7 @@ func FTAggregateQuery(query string, options *FTAggregateOptions) AggregateQuery
571585
for _, sortBy := range options.SortBy {
572586
sortByOptions = append(sortByOptions, sortBy.FieldName)
573587
if sortBy.Asc && sortBy.Desc {
574-
panic("FT.AGGREGATE: ASC and DESC are mutually exclusive")
588+
return nil, fmt.Errorf("FT.AGGREGATE: ASC and DESC are mutually exclusive")
575589
}
576590
if sortBy.Asc {
577591
sortByOptions = append(sortByOptions, "ASC")
@@ -616,7 +630,7 @@ func FTAggregateQuery(query string, options *FTAggregateOptions) AggregateQuery
616630
queryArgs = append(queryArgs, "DIALECT", 2)
617631
}
618632
}
619-
return queryArgs
633+
return queryArgs, nil
620634
}
621635

622636
func ProcessAggregateResult(data []interface{}) (*FTAggregateResult, error) {
@@ -718,7 +732,9 @@ func (c cmdable) FTAggregateWithArgs(ctx context.Context, index string, query st
718732
args = append(args, "ADDSCORES")
719733
}
720734
if options.LoadAll && options.Load != nil {
721-
panic("FT.AGGREGATE: LOADALL and LOAD are mutually exclusive")
735+
cmd := NewAggregateCmd(ctx, args...)
736+
cmd.SetErr(fmt.Errorf("FT.AGGREGATE: LOADALL and LOAD are mutually exclusive"))
737+
return cmd
722738
}
723739
if options.LoadAll {
724740
args = append(args, "LOAD", "*")
@@ -771,7 +787,9 @@ func (c cmdable) FTAggregateWithArgs(ctx context.Context, index string, query st
771787
for _, sortBy := range options.SortBy {
772788
sortByOptions = append(sortByOptions, sortBy.FieldName)
773789
if sortBy.Asc && sortBy.Desc {
774-
panic("FT.AGGREGATE: ASC and DESC are mutually exclusive")
790+
cmd := NewAggregateCmd(ctx, args...)
791+
cmd.SetErr(fmt.Errorf("FT.AGGREGATE: ASC and DESC are mutually exclusive"))
792+
return cmd
775793
}
776794
if sortBy.Asc {
777795
sortByOptions = append(sortByOptions, "ASC")
@@ -919,7 +937,9 @@ func (c cmdable) FTCreate(ctx context.Context, index string, options *FTCreateOp
919937
args = append(args, "ON", "JSON")
920938
}
921939
if options.OnHash && options.OnJSON {
922-
panic("FT.CREATE: ON HASH and ON JSON are mutually exclusive")
940+
cmd := NewStatusCmd(ctx, args...)
941+
cmd.SetErr(fmt.Errorf("FT.CREATE: ON HASH and ON JSON are mutually exclusive"))
942+
return cmd
923943
}
924944
if options.Prefix != nil {
925945
args = append(args, "PREFIX", len(options.Prefix))
@@ -970,12 +990,16 @@ func (c cmdable) FTCreate(ctx context.Context, index string, options *FTCreateOp
970990
}
971991
}
972992
if schema == nil {
973-
panic("FT.CREATE: SCHEMA is required")
993+
cmd := NewStatusCmd(ctx, args...)
994+
cmd.SetErr(fmt.Errorf("FT.CREATE: SCHEMA is required"))
995+
return cmd
974996
}
975997
args = append(args, "SCHEMA")
976998
for _, schema := range schema {
977999
if schema.FieldName == "" || schema.FieldType == SearchFieldTypeInvalid {
978-
panic("FT.CREATE: SCHEMA FieldName and FieldType are required")
1000+
cmd := NewStatusCmd(ctx, args...)
1001+
cmd.SetErr(fmt.Errorf("FT.CREATE: SCHEMA FieldName and FieldType are required"))
1002+
return cmd
9791003
}
9801004
args = append(args, schema.FieldName)
9811005
if schema.As != "" {
@@ -984,15 +1008,32 @@ func (c cmdable) FTCreate(ctx context.Context, index string, options *FTCreateOp
9841008
args = append(args, schema.FieldType.String())
9851009
if schema.VectorArgs != nil {
9861010
if schema.FieldType != SearchFieldTypeVector {
987-
panic("FT.CREATE: SCHEMA FieldType VECTOR is required for VectorArgs")
1011+
cmd := NewStatusCmd(ctx, args...)
1012+
cmd.SetErr(fmt.Errorf("FT.CREATE: SCHEMA FieldType VECTOR is required for VectorArgs"))
1013+
return cmd
1014+
}
1015+
// Check mutual exclusivity of vector options
1016+
optionCount := 0
1017+
if schema.VectorArgs.FlatOptions != nil {
1018+
optionCount++
1019+
}
1020+
if schema.VectorArgs.HNSWOptions != nil {
1021+
optionCount++
1022+
}
1023+
if schema.VectorArgs.VamanaOptions != nil {
1024+
optionCount++
9881025
}
989-
if schema.VectorArgs.FlatOptions != nil && schema.VectorArgs.HNSWOptions != nil {
990-
panic("FT.CREATE: SCHEMA VectorArgs FlatOptions and HNSWOptions are mutually exclusive")
1026+
if optionCount != 1 {
1027+
cmd := NewStatusCmd(ctx, args...)
1028+
cmd.SetErr(fmt.Errorf("FT.CREATE: SCHEMA VectorArgs must have exactly one of FlatOptions, HNSWOptions, or VamanaOptions"))
1029+
return cmd
9911030
}
9921031
if schema.VectorArgs.FlatOptions != nil {
9931032
args = append(args, "FLAT")
9941033
if schema.VectorArgs.FlatOptions.Type == "" || schema.VectorArgs.FlatOptions.Dim == 0 || schema.VectorArgs.FlatOptions.DistanceMetric == "" {
995-
panic("FT.CREATE: Type, Dim and DistanceMetric are required for VECTOR FLAT")
1034+
cmd := NewStatusCmd(ctx, args...)
1035+
cmd.SetErr(fmt.Errorf("FT.CREATE: Type, Dim and DistanceMetric are required for VECTOR FLAT"))
1036+
return cmd
9961037
}
9971038
flatArgs := []interface{}{
9981039
"TYPE", schema.VectorArgs.FlatOptions.Type,
@@ -1011,7 +1052,9 @@ func (c cmdable) FTCreate(ctx context.Context, index string, options *FTCreateOp
10111052
if schema.VectorArgs.HNSWOptions != nil {
10121053
args = append(args, "HNSW")
10131054
if schema.VectorArgs.HNSWOptions.Type == "" || schema.VectorArgs.HNSWOptions.Dim == 0 || schema.VectorArgs.HNSWOptions.DistanceMetric == "" {
1014-
panic("FT.CREATE: Type, Dim and DistanceMetric are required for VECTOR HNSW")
1055+
cmd := NewStatusCmd(ctx, args...)
1056+
cmd.SetErr(fmt.Errorf("FT.CREATE: Type, Dim and DistanceMetric are required for VECTOR HNSW"))
1057+
return cmd
10151058
}
10161059
hnswArgs := []interface{}{
10171060
"TYPE", schema.VectorArgs.HNSWOptions.Type,
@@ -1036,10 +1079,48 @@ func (c cmdable) FTCreate(ctx context.Context, index string, options *FTCreateOp
10361079
args = append(args, len(hnswArgs))
10371080
args = append(args, hnswArgs...)
10381081
}
1082+
if schema.VectorArgs.VamanaOptions != nil {
1083+
args = append(args, "SVS-VAMANA")
1084+
if schema.VectorArgs.VamanaOptions.Type == "" || schema.VectorArgs.VamanaOptions.Dim == 0 || schema.VectorArgs.VamanaOptions.DistanceMetric == "" {
1085+
cmd := NewStatusCmd(ctx, args...)
1086+
cmd.SetErr(fmt.Errorf("FT.CREATE: Type, Dim and DistanceMetric are required for VECTOR VAMANA"))
1087+
return cmd
1088+
}
1089+
vamanaArgs := []interface{}{
1090+
"TYPE", schema.VectorArgs.VamanaOptions.Type,
1091+
"DIM", schema.VectorArgs.VamanaOptions.Dim,
1092+
"DISTANCE_METRIC", schema.VectorArgs.VamanaOptions.DistanceMetric,
1093+
}
1094+
if schema.VectorArgs.VamanaOptions.Compression != "" {
1095+
vamanaArgs = append(vamanaArgs, "COMPRESSION", schema.VectorArgs.VamanaOptions.Compression)
1096+
}
1097+
if schema.VectorArgs.VamanaOptions.ConstructionWindowSize > 0 {
1098+
vamanaArgs = append(vamanaArgs, "CONSTRUCTION_WINDOW_SIZE", schema.VectorArgs.VamanaOptions.ConstructionWindowSize)
1099+
}
1100+
if schema.VectorArgs.VamanaOptions.GraphMaxDegree > 0 {
1101+
vamanaArgs = append(vamanaArgs, "GRAPH_MAX_DEGREE", schema.VectorArgs.VamanaOptions.GraphMaxDegree)
1102+
}
1103+
if schema.VectorArgs.VamanaOptions.SearchWindowSize > 0 {
1104+
vamanaArgs = append(vamanaArgs, "SEARCH_WINDOW_SIZE", schema.VectorArgs.VamanaOptions.SearchWindowSize)
1105+
}
1106+
if schema.VectorArgs.VamanaOptions.Epsilon > 0 {
1107+
vamanaArgs = append(vamanaArgs, "EPSILON", schema.VectorArgs.VamanaOptions.Epsilon)
1108+
}
1109+
if schema.VectorArgs.VamanaOptions.TrainingThreshold > 0 {
1110+
vamanaArgs = append(vamanaArgs, "TRAINING_THRESHOLD", schema.VectorArgs.VamanaOptions.TrainingThreshold)
1111+
}
1112+
if schema.VectorArgs.VamanaOptions.ReduceDim > 0 {
1113+
vamanaArgs = append(vamanaArgs, "REDUCE", schema.VectorArgs.VamanaOptions.ReduceDim)
1114+
}
1115+
args = append(args, len(vamanaArgs))
1116+
args = append(args, vamanaArgs...)
1117+
}
10391118
}
10401119
if schema.GeoShapeFieldType != "" {
10411120
if schema.FieldType != SearchFieldTypeGeoShape {
1042-
panic("FT.CREATE: SCHEMA FieldType GEOSHAPE is required for GeoShapeFieldType")
1121+
cmd := NewStatusCmd(ctx, args...)
1122+
cmd.SetErr(fmt.Errorf("FT.CREATE: SCHEMA FieldType GEOSHAPE is required for GeoShapeFieldType"))
1123+
return cmd
10431124
}
10441125
args = append(args, schema.GeoShapeFieldType)
10451126
}
@@ -1197,7 +1278,7 @@ func (c cmdable) FTExplainWithArgs(ctx context.Context, index string, query stri
11971278
// FTExplainCli - Returns the execution plan for a complex query. [Not Implemented]
11981279
// For more information, see https://redis.io/commands/ft.explaincli/
11991280
func (c cmdable) FTExplainCli(ctx context.Context, key, path string) error {
1200-
panic("not implemented")
1281+
return fmt.Errorf("FTExplainCli is not implemented")
12011282
}
12021283

12031284
func parseFTInfo(data map[string]interface{}) (FTInfoResult, error) {
@@ -1758,7 +1839,7 @@ type SearchQuery []interface{}
17581839
// For more information, please refer to the Redis documentation about [FT.SEARCH].
17591840
//
17601841
// [FT.SEARCH]: (https://redis.io/commands/ft.search/)
1761-
func FTSearchQuery(query string, options *FTSearchOptions) SearchQuery {
1842+
func FTSearchQuery(query string, options *FTSearchOptions) (SearchQuery, error) {
17621843
queryArgs := []interface{}{query}
17631844
if options != nil {
17641845
if options.NoContent {
@@ -1838,7 +1919,7 @@ func FTSearchQuery(query string, options *FTSearchOptions) SearchQuery {
18381919
for _, sortBy := range options.SortBy {
18391920
queryArgs = append(queryArgs, sortBy.FieldName)
18401921
if sortBy.Asc && sortBy.Desc {
1841-
panic("FT.SEARCH: ASC and DESC are mutually exclusive")
1922+
return nil, fmt.Errorf("FT.SEARCH: ASC and DESC are mutually exclusive")
18421923
}
18431924
if sortBy.Asc {
18441925
queryArgs = append(queryArgs, "ASC")
@@ -1866,7 +1947,7 @@ func FTSearchQuery(query string, options *FTSearchOptions) SearchQuery {
18661947
queryArgs = append(queryArgs, "DIALECT", 2)
18671948
}
18681949
}
1869-
return queryArgs
1950+
return queryArgs, nil
18701951
}
18711952

18721953
// FTSearchWithArgs - Executes a search query on an index with additional options.
@@ -1955,7 +2036,9 @@ func (c cmdable) FTSearchWithArgs(ctx context.Context, index string, query strin
19552036
for _, sortBy := range options.SortBy {
19562037
args = append(args, sortBy.FieldName)
19572038
if sortBy.Asc && sortBy.Desc {
1958-
panic("FT.SEARCH: ASC and DESC are mutually exclusive")
2039+
cmd := newFTSearchCmd(ctx, options, args...)
2040+
cmd.SetErr(fmt.Errorf("FT.SEARCH: ASC and DESC are mutually exclusive"))
2041+
return cmd
19592042
}
19602043
if sortBy.Asc {
19612044
args = append(args, "ASC")

0 commit comments

Comments
 (0)