Support Vector Similarity (#151)

* Support Vector Similarity

Co-authored-by: Chayim I. Kirshen <[email protected]>

Author: Avital-Fine

.circleci/config.yml

@@ -53,7 +53,7 @@ jobs:
   build: # test with redisearch:latest
     docker:
       - image: circleci/golang:1.16
-      - image: redislabs/redisearch:latest
+      - image: redislabs/redisearch:edge
 
     working_directory: /go/src/github.com/RediSearch/redisearch-go
     steps:
@@ -65,34 +65,12 @@ jobs:
       - run: make coverage
       - run: bash <(curl -s https://raw.githubusercontent.com/codecov/codecov-bash/master/codecov) -t ${CODECOV_TOKEN}
 
-  build-v16:
-    docker:
-      - image: circleci/golang:1.16
-      - image: redislabs/redisearch:1.6.15
-
-    working_directory: /go/src/github.com/RediSearch/redisearch-go
-    steps:
-      - checkout
-      - run: make test
-
-  build-edge: # test nightly with redisearch:edge
-    docker:
-      - image: circleci/golang:1.16
-      - image: redislabs/redisearch:edge
-
-    working_directory: /go/src/github.com/RediSearch/redisearch-go
-    steps:
-      - checkout
-      - run: make test
-
 workflows:
   version: 2
   commit:
     jobs:
       - build-tls
       - build
-      - build-edge
-      - build-v16
   nightly:
     triggers:
       - schedule:
@@ -102,5 +80,5 @@ workflows:
               only:
                 - master
     jobs:
-      - build-edge
+      - build
       - build-tls

redisearch/query.go

@@ -93,6 +93,8 @@ type Query struct {
 	SortBy        *SortingKey
 	HighlightOpts *HighlightOptions
 	SummarizeOpts *SummaryOptions
+	Params        map[string]interface{}
+	Dialect       int
 }
 
 // Paging represents the offset paging of a search result
@@ -227,6 +229,18 @@ func (q Query) serialize() redis.Args {
 			}
 		}
 	}
+
+	if q.Params != nil {
+		args = args.Add("PARAMS", len(q.Params)*2)
+		for name, value := range q.Params {
+			args = args.Add(name, value)
+		}
+	}
+
+	if q.Dialect != 0 {
+		args = args.Add("DIALECT", q.Dialect)
+	}
+
 	return args
 }
 
@@ -370,6 +384,29 @@ func (q *Query) SummarizeOptions(opts SummaryOptions) *Query {
 	return q
 }
 
+// SetParams sets parameters that can be referenced in the query string by a $ , followed by the parameter name,
+// e.g., $user , and each such reference in the search query to a parameter name is substituted
+// by the corresponding parameter value.
+func (q *Query) SetParams(params map[string]interface{}) *Query {
+	q.Params = params
+	return q
+}
+
+// AddParam adds a new param to the parameters list
+func (q *Query) AddParam(name string, value interface{}) *Query {
+	if q.Params == nil {
+		q.Params = make(map[string]interface{})
+	}
+	q.Params[name] = value
+	return q
+}
+
+// SetDialect can have one of 2 options: 1 or 2
+func (q *Query) SetDialect(dialect int) *Query {
+	q.Dialect = dialect
+	return q
+}
+
 // IndexOptions indexes multiple documents on the index, with optional Options passed to options
 func (i *Client) IndexOptions(opts IndexingOptions, docs ...Document) error {
 

redisearch/query_test.go

@@ -79,6 +79,8 @@ func TestQuery_serialize(t *testing.T) {
 		SortBy        *SortingKey
 		HighlightOpts *HighlightOptions
 		SummarizeOpts *SummaryOptions
+		Params        map[string]interface{}
+		Dialect       int
 	}
 	tests := []struct {
 		name   string
@@ -111,6 +113,8 @@ func TestQuery_serialize(t *testing.T) {
 			NumFragments: 3,
 			Separator:    "...",
 		}}, redis.Args{raw, "LIMIT", 0, 0, "SUMMARIZE", "FIELDS", 1, "test_field", "LEN", 20, "FRAGS", 3, "SEPARATOR", "..."}},
+		{"Params", fields{Raw: raw, Params: map[string]interface{}{"min": 1}}, redis.Args{raw, "LIMIT", 0, 0, "PARAMS", 2, "min", 1}},
+		{"Dialect", fields{Raw: raw, Dialect: 2}, redis.Args{raw, "LIMIT", 0, 0, "DIALECT", 2}},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -126,6 +130,8 @@ func TestQuery_serialize(t *testing.T) {
 				SortBy:        tt.fields.SortBy,
 				HighlightOpts: tt.fields.HighlightOpts,
 				SummarizeOpts: tt.fields.SummarizeOpts,
+				Params:        tt.fields.Params,
+				Dialect:       tt.fields.Dialect,
 			}
 			if g := q.serialize(); !reflect.DeepEqual(g, tt.want) {
 				t.Errorf("serialize() = %v, want %v", g, tt.want)

redisearch/redisearch_test.go

@@ -510,3 +510,67 @@ func TestReturnFields(t *testing.T) {
 	assert.Equal(t, "Jon", docs[0].Properties["name"])
 	assert.Equal(t, "25", docs[0].Properties["years"])
 }
+
+func TestParams(t *testing.T) {
+	c := createClient("TestParams")
+	version, _ := c.getRediSearchVersion()
+	if version < 20430 {
+		// VectorSimilarity is available for RediSearch 2.2+
+		return
+	}
+
+	// Create a schema
+	sc := NewSchema(DefaultOptions).AddField(NewNumericField("numval"))
+	c.Drop()
+	assert.Nil(t, c.CreateIndex(sc))
+	// Create data
+	_, err := c.pool.Get().Do("HSET", "1", "numval", "1")
+	assert.Nil(t, err)
+	_, err = c.pool.Get().Do("HSET", "2", "numval", "2")
+	assert.Nil(t, err)
+	_, err = c.pool.Get().Do("HSET", "3", "numval", "3")
+	assert.Nil(t, err)
+	// Searching with parameters
+	_, total, err := c.Search(NewQuery("@numval:[$min $max]").
+		SetParams(map[string]interface{}{"min": "1", "max": "2"}).
+		SetDialect(2))
+	assert.Nil(t, err)
+	assert.Equal(t, 2, total)
+}
+
+func TestVectorField(t *testing.T) {
+	c := createClient("TestVectorField")
+	version, _ := c.getRediSearchVersion()
+	if version < 20430 {
+		// VectorSimilarity is available for RediSearch 2.2+
+		return
+	}
+
+	// Create a schema
+	sc := NewSchema(DefaultOptions).AddField(
+		NewVectorFieldOptions("v", VectorFieldOptions{Algorithm: Flat, Attributes: map[string]interface{}{
+			"TYPE":            "FLOAT32",
+			"DIM":             2,
+			"DISTANCE_METRIC": "L2",
+		}}),
+	)
+	c.Drop()
+	assert.Nil(t, c.CreateIndex(sc))
+	// Create data
+	_, err := c.pool.Get().Do("HSET", "a", "v", "aaaaaaaa")
+	assert.Nil(t, err)
+	_, err = c.pool.Get().Do("HSET", "b", "v", "aaaabaaa")
+	assert.Nil(t, err)
+	_, err = c.pool.Get().Do("HSET", "c", "v", "aaaaabaa")
+	assert.Nil(t, err)
+	// Searching with parameters
+	docs, total, err := c.Search(NewQuery("*=>[KNN 2 @v $vec]").
+		AddParam("vec", "aaaaaaaa").
+		SetSortBy("__v_score", true).
+		AddReturnFields("__v_score").
+		SetDialect(2))
+	assert.Nil(t, err)
+	assert.Equal(t, 2, total)
+	assert.Equal(t, "a", docs[0].Id)
+	assert.Equal(t, "0", docs[0].Properties["__v_score"])
+}

redisearch/schema.go

@@ -135,6 +135,9 @@ const (
 
 	// TagField is a field used for compact indexing of comma separated values
 	TagField
+
+	//VectorField allows vector similarity queries against the value in this attribute.
+	VectorField
 )
 
 // Phonetic Matchers
@@ -185,6 +188,20 @@ type GeoFieldOptions struct {
 	As      string
 }
 
+type algorithm string
+
+// Supported algorithms for Vector field
+const (
+	Flat algorithm = "FLAT"
+	HNSW algorithm = "HNSW"
+)
+
+// VectorFieldOptions Options for vector fields
+type VectorFieldOptions struct {
+	Algorithm  algorithm
+	Attributes map[string]interface{}
+}
+
 // NewTextField creates a new text field with the given weight
 func NewTextField(name string) Field {
 	return Field{
@@ -268,6 +285,15 @@ func NewGeoFieldOptions(name string, options GeoFieldOptions) Field {
 	return f
 }
 
+// NewVectorFieldOptions creates a new geo field with the given name and additional options
+func NewVectorFieldOptions(name string, options VectorFieldOptions) Field {
+	return Field{
+		Name:    name,
+		Type:    VectorField,
+		Options: options,
+	}
+}
+
 // Schema represents an index schema Schema, or how the index would
 // treat documents sent to it.
 type Schema struct {
@@ -417,6 +443,26 @@ func serializeField(f Field, args redis.Args) (argsOut redis.Args, err error) {
 				argsOut = append(argsOut, "NOINDEX")
 			}
 		}
+	case VectorField:
+		argsOut = append(argsOut, f.Name, "VECTOR")
+		if f.Options != nil {
+			opts, ok := f.Options.(VectorFieldOptions)
+			if !ok {
+				err = fmt.Errorf("Error on VectorField serialization")
+				return
+			}
+			if opts.Algorithm != "" {
+				argsOut = append(argsOut, opts.Algorithm)
+			}
+			if opts.Attributes != nil {
+				var flat []interface{}
+				for attrName, attrValue := range opts.Attributes {
+					flat = append(flat, attrName, attrValue)
+				}
+				argsOut = append(argsOut, len(flat))
+				argsOut = append(argsOut, flat...)
+			}
+		}
 	default:
 		err = fmt.Errorf("Unrecognized field type %v serialization", f.Type)
 		return

redisearch/schema_test.go

@@ -74,6 +74,7 @@ func TestSerializeSchema(t *testing.T) {
 		{"default-and-tag", args{NewSchema(DefaultOptions).AddField(NewTagField("tag-field")), redis.Args{}}, redis.Args{"SCHEMA", "tag-field", "TAG", "SEPARATOR", ","}, false},
 		{"default-and-tag-with-options", args{NewSchema(DefaultOptions).AddField(NewTagFieldOptions("tag-field", TagFieldOptions{Sortable: true, NoIndex: false, Separator: byte(','), As: "field"})), redis.Args{}}, redis.Args{"SCHEMA", "tag-field", "AS", "field", "TAG", "SEPARATOR", ",", "SORTABLE"}, false},
 		{"default-geo-with-options", args{NewSchema(DefaultOptions).AddField(NewGeoFieldOptions("location", GeoFieldOptions{As: "loc"})), redis.Args{}}, redis.Args{"SCHEMA", "location", "AS", "loc", "GEO"}, false},
+		{"default-vector", args{NewSchema(DefaultOptions).AddField(NewVectorFieldOptions("vec", VectorFieldOptions{Algorithm: Flat, Attributes: map[string]interface{}{"DIM": 128}})), redis.Args{}}, redis.Args{"SCHEMA", "vec", "VECTOR", Flat, 2, "DIM", 128}, false},
 		{"error-unsupported", args{NewSchema(DefaultOptions).AddField(Field{Type: 10}), redis.Args{}}, nil, true},
 	}
 	for _, tt := range tests {