Boost Suggestions by Context

suggester_completion.go

@@ -123,7 +123,7 @@ func (q *CompletionSuggester) Source(includeName bool) (interface{}, error) {
 				ctxq[k] = v
 			}
 		}
-		suggester["context"] = ctxq
+		suggester["contexts"] = ctxq
 	}
 
 	// TODO(oe) Add completion-suggester specific parameters here

suggester_completion_test.go

@@ -45,8 +45,8 @@ func TestCompletionSuggesterSourceWithMultipleContexts(t *testing.T) {
 		t.Fatalf("marshaling to JSON failed: %v", err)
 	}
 	got := string(data)
-	expected := `{"song-suggest":{"text":"n","completion":{"context":{"artist":"Sting","label":"BMG"},"field":"suggest"}}}`
+	expected := `{"song-suggest":{"text":"n","completion":{"contexts":{"artist":[{"context":"Sting"}],"label":[{"context":"BMG"}]},"field":"suggest"}}}`
 	if got != expected {
-		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
+		t.Errorf("expected %s\n,got:\n%s", expected, got)
 	}
 }

suggester_context_category.go

@@ -62,23 +62,36 @@ func (q *SuggesterCategoryMapping) Source() (interface{}, error) {
 // See https://www.elastic.co/guide/en/elasticsearch/reference/5.2/suggester-context.html#_category_query.
 type SuggesterCategoryQuery struct {
 	name   string
-	values []string
+	values map[string]*int
 }
 
 // NewSuggesterCategoryQuery creates a new SuggesterCategoryQuery.
 func NewSuggesterCategoryQuery(name string, values ...string) *SuggesterCategoryQuery {
 	q := &SuggesterCategoryQuery{
 		name:   name,
-		values: make([]string, 0),
+		values: make(map[string]*int),
 	}
+
 	if len(values) > 0 {
-		q.values = append(q.values, values...)
+		q.Values(values...)
 	}
 	return q
 }
 
+func (q *SuggesterCategoryQuery) Value(val string) *SuggesterCategoryQuery {
+	q.values[val] = nil
+	return q
+}
+
+func (q *SuggesterCategoryQuery) ValueWithBoost(val string, boost int) *SuggesterCategoryQuery {
+	q.values[val] = &boost
+	return q
+}
+
 func (q *SuggesterCategoryQuery) Values(values ...string) *SuggesterCategoryQuery {
-	q.values = append(q.values, values...)
+	for _, val := range values {
+		q.values[val] = nil
+	}
 	return q
 }
 
@@ -88,11 +101,18 @@ func (q *SuggesterCategoryQuery) Source() (interface{}, error) {
 
 	switch len(q.values) {
 	case 0:
-		source[q.name] = q.values
-	case 1:
-		source[q.name] = q.values[0]
+		source[q.name] = make([]string, 0)
 	default:
-		source[q.name] = q.values
+		contexts := make([]interface{}, 0)
+		for val, boost := range q.values {
+			context := make(map[string]interface{})
+			context["context"] = val
+			if boost != nil {
+				context["boost"] = *boost
+			}
+			contexts = append(contexts, context)
+		}
+		source[q.name] = contexts
 	}
 
 	return source, nil

suggester_context_category_test.go

@@ -73,7 +73,7 @@ func TestSuggesterCategoryQuery(t *testing.T) {
 		t.Fatalf("marshaling to JSON failed: %v", err)
 	}
 	got := string(data)
-	expected := `{"color":"red"}`
+	expected := `{"color":[{"context":"red"}]}`
 	if got != expected {
 		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
 	}
@@ -90,8 +90,44 @@ func TestSuggesterCategoryQueryWithTwoValues(t *testing.T) {
 		t.Fatalf("marshaling to JSON failed: %v", err)
 	}
 	got := string(data)
-	expected := `{"color":["red","yellow"]}`
+	expected := `{"color":[{"context":"red"},{"context":"yellow"}]}`
 	if got != expected {
-		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
+		t.Errorf("expected %s\n,got:\n%s", expected, got)
+	}
+}
+
+func TestSuggesterCategoryQueryWithBoost(t *testing.T) {
+	q := NewSuggesterCategoryQuery("color", "red")
+	q.ValueWithBoost("yellow", 4)
+	src, err := q.Source()
+	if err != nil {
+		t.Fatal(err)
+	}
+	data, err := json.Marshal(src)
+	if err != nil {
+		t.Fatalf("marshaling to JSON failed: %v", err)
+	}
+	got := string(data)
+	expected := `{"color":[{"context":"red"},{"boost":4,"context":"yellow"}]}`
+	if got != expected {
+		t.Errorf("expected %s\n,got:\n%s", expected, got)
+	}
+}
+
+func TestSuggesterCategoryQueryWithoutBoost(t *testing.T) {
+	q := NewSuggesterCategoryQuery("color", "red")
+	q.Value("yellow")
+	src, err := q.Source()
+	if err != nil {
+		t.Fatal(err)
+	}
+	data, err := json.Marshal(src)
+	if err != nil {
+		t.Fatalf("marshaling to JSON failed: %v", err)
+	}
+	got := string(data)
+	expected := `{"color":[{"context":"red"},{"context":"yellow"}]}`
+	if got != expected {
+		t.Errorf("expected %s\n,got:\n%s", expected, got)
 	}
 }