Skip to content

Commit

Permalink
Add RankFeatureQuery
Browse files Browse the repository at this point in the history
The `rank_feature` query boosts the relevance score of documents based
on the numeric value of a `rank_feature` or `rank_features` field.

See here for details:
https://www.elastic.co/guide/en/elasticsearch/reference/7.14/query-dsl-rank-feature-query.html

Close #1428
  • Loading branch information
olivere committed Sep 22, 2021
1 parent 593cd32 commit 987f2d7
Show file tree
Hide file tree
Showing 3 changed files with 284 additions and 0 deletions.
4 changes: 4 additions & 0 deletions search_queries_pinned_test.go
@@ -1,3 +1,7 @@
// Copyright 2012-present Oliver Eilhard. All rights reserved.
// Use of this source code is governed by a MIT-license.
// See http://olivere.mit-license.org/license.txt for details.

package elastic

import (
Expand Down
213 changes: 213 additions & 0 deletions search_queries_rank_feature.go
@@ -0,0 +1,213 @@
// Copyright 2012-present Oliver Eilhard. All rights reserved.
// Use of this source code is governed by a MIT-license.
// See http://olivere.mit-license.org/license.txt for details.

package elastic

// RankFeatureQuery boosts the relevance score of documents based on the
// numeric value of a rank_feature or rank_features field.
//
// The RankFeatureQuery is typically used in the should clause of a BoolQuery
// so its relevance scores are added to other scores from the BoolQuery.
//
// For more details, see:
// https://www.elastic.co/guide/en/elasticsearch/reference/7.14/query-dsl-rank-feature-query.html
type RankFeatureQuery struct {
field string
scoreFunc RankFeatureScoreFunction
boost *float64
queryName string
}

// NewRankFeatureQuery creates and initializes a new RankFeatureQuery.
func NewRankFeatureQuery(field string) *RankFeatureQuery {
return &RankFeatureQuery{
field: field,
}
}

// Field name.
func (q *RankFeatureQuery) Field(field string) *RankFeatureQuery {
q.field = field
return q
}

// ScoreFunction specifies the score function for the RankFeatureQuery.
func (q *RankFeatureQuery) ScoreFunction(f RankFeatureScoreFunction) *RankFeatureQuery {
q.scoreFunc = f
return q
}

// Boost sets the boost for this query.
func (q *RankFeatureQuery) Boost(boost float64) *RankFeatureQuery {
q.boost = &boost
return q
}

// QueryName sets the query name for the filter that can be used when
// searching for matched_filters per hit.
func (q *RankFeatureQuery) QueryName(queryName string) *RankFeatureQuery {
q.queryName = queryName
return q
}

// Source returns the JSON serializable content for this query.
func (q *RankFeatureQuery) Source() (interface{}, error) {
// {
// "rank_feature": {
// "field": "pagerank",
// "saturation": {
// "pivot": 8
// }
// }
// }

query := make(map[string]interface{})
params := make(map[string]interface{})
query["rank_feature"] = params
params["field"] = q.field
if q.scoreFunc != nil {
src, err := q.scoreFunc.Source()
if err != nil {
return nil, err
}
params[q.scoreFunc.Name()] = src
}
if q.boost != nil {
params["boost"] = *q.boost
}
if q.queryName != "" {
params["_name"] = q.queryName
}

return query, nil
}

// -- Score functions --

// RankFeatureScoreFunction specifies the interface for score functions
// in the context of a RankFeatureQuery.
type RankFeatureScoreFunction interface {
Name() string
Source() (interface{}, error)
}

// -- Log score function --

// RankFeatureLogScoreFunction represents a Logarithmic score function for a
// RankFeatureQuery.
//
// See here for details:
// https://www.elastic.co/guide/en/elasticsearch/reference/7.14/query-dsl-rank-feature-query.html#rank-feature-query-logarithm
type RankFeatureLogScoreFunction struct {
scalingFactor float64
}

// NewRankFeatureLogScoreFunction returns a new RankFeatureLogScoreFunction
// with the given scaling factor.
func NewRankFeatureLogScoreFunction(scalingFactor float64) *RankFeatureLogScoreFunction {
return &RankFeatureLogScoreFunction{
scalingFactor: scalingFactor,
}
}

// Name of the score function.
func (f *RankFeatureLogScoreFunction) Name() string { return "log" }

// Source returns a serializable JSON object for building the query.
func (f *RankFeatureLogScoreFunction) Source() (interface{}, error) {
return map[string]interface{}{
"scaling_factor": f.scalingFactor,
}, nil
}

// -- Saturation score function --

// RankFeatureSaturationScoreFunction represents a Log score function for a
// RankFeatureQuery.
//
// See here for details:
// https://www.elastic.co/guide/en/elasticsearch/reference/7.14/query-dsl-rank-feature-query.html#rank-feature-query-saturation
type RankFeatureSaturationScoreFunction struct {
pivot *float64
}

// NewRankFeatureSaturationScoreFunction initializes a new
// RankFeatureSaturationScoreFunction.
func NewRankFeatureSaturationScoreFunction() *RankFeatureSaturationScoreFunction {
return &RankFeatureSaturationScoreFunction{}
}

// Pivot specifies the pivot to use.
func (f *RankFeatureSaturationScoreFunction) Pivot(pivot float64) *RankFeatureSaturationScoreFunction {
f.pivot = &pivot
return f
}

// Name of the score function.
func (f *RankFeatureSaturationScoreFunction) Name() string { return "saturation" }

// Source returns a serializable JSON object for building the query.
func (f *RankFeatureSaturationScoreFunction) Source() (interface{}, error) {
m := make(map[string]interface{})
if f.pivot != nil {
m["pivot"] = *f.pivot
}
return m, nil
}

// -- Sigmoid score function --

// RankFeatureSigmoidScoreFunction represents a Sigmoid score function for a
// RankFeatureQuery.
//
// See here for details:
// https://www.elastic.co/guide/en/elasticsearch/reference/7.14/query-dsl-rank-feature-query.html#rank-feature-query-sigmoid
type RankFeatureSigmoidScoreFunction struct {
pivot float64
exponent float64
}

// NewRankFeatureSigmoidScoreFunction returns a new RankFeatureSigmoidScoreFunction
// with the given scaling factor.
func NewRankFeatureSigmoidScoreFunction(pivot, exponent float64) *RankFeatureSigmoidScoreFunction {
return &RankFeatureSigmoidScoreFunction{
pivot: pivot,
exponent: exponent,
}
}

// Name of the score function.
func (f *RankFeatureSigmoidScoreFunction) Name() string { return "sigmoid" }

// Source returns a serializable JSON object for building the query.
func (f *RankFeatureSigmoidScoreFunction) Source() (interface{}, error) {
return map[string]interface{}{
"pivot": f.pivot,
"exponent": f.exponent,
}, nil
}

// -- Linear score function --

// RankFeatureLinearScoreFunction represents a Linear score function for a
// RankFeatureQuery.
//
// See here for details:
// https://www.elastic.co/guide/en/elasticsearch/reference/7.14/query-dsl-rank-feature-query.html#rank-feature-query-linear
type RankFeatureLinearScoreFunction struct {
}

// NewRankFeatureLinearScoreFunction initializes a new
// RankFeatureLinearScoreFunction.
func NewRankFeatureLinearScoreFunction() *RankFeatureLinearScoreFunction {
return &RankFeatureLinearScoreFunction{}
}

// Name of the score function.
func (f *RankFeatureLinearScoreFunction) Name() string { return "linear" }

// Source returns a serializable JSON object for building the query.
func (f *RankFeatureLinearScoreFunction) Source() (interface{}, error) {
return map[string]interface{}{}, nil
}
67 changes: 67 additions & 0 deletions search_queries_rank_feature_test.go
@@ -0,0 +1,67 @@
// Copyright 2012-present Oliver Eilhard. All rights reserved.
// Use of this source code is governed by a MIT-license.
// See http://olivere.mit-license.org/license.txt for details.

package elastic

import (
"encoding/json"
"testing"
)

func TestRankFeatureQueryTest(t *testing.T) {
tests := []struct {
Query Query
Expected string
}{
// #0
{
Query: NewRankFeatureQuery("pagerank"),
Expected: `{"rank_feature":{"field":"pagerank"}}`,
},
// #1
{
Query: NewRankFeatureQuery("url_length").Boost(0.1),
Expected: `{"rank_feature":{"boost":0.1,"field":"url_length"}}`,
},
// #2
{
Query: NewRankFeatureQuery("pagerank").ScoreFunction(NewRankFeatureSaturationScoreFunction().Pivot(8)),
Expected: `{"rank_feature":{"field":"pagerank","saturation":{"pivot":8}}}`,
},
// #3
{
Query: NewRankFeatureQuery("pagerank").ScoreFunction(NewRankFeatureSaturationScoreFunction()),
Expected: `{"rank_feature":{"field":"pagerank","saturation":{}}}`,
},
// #4
{
Query: NewRankFeatureQuery("pagerank").ScoreFunction(NewRankFeatureLogScoreFunction(4)),
Expected: `{"rank_feature":{"field":"pagerank","log":{"scaling_factor":4}}}`,
},
// #5
{
Query: NewRankFeatureQuery("pagerank").ScoreFunction(NewRankFeatureSigmoidScoreFunction(7, 0.6)),
Expected: `{"rank_feature":{"field":"pagerank","sigmoid":{"exponent":0.6,"pivot":7}}}`,
},
// #6
{
Query: NewRankFeatureQuery("pagerank").ScoreFunction(NewRankFeatureLinearScoreFunction()),
Expected: `{"rank_feature":{"field":"pagerank","linear":{}}}`,
},
}

for i, tt := range tests {
src, err := tt.Query.Source()
if err != nil {
t.Fatalf("#%d: encoding Source failed: %v", i, err)
}
data, err := json.Marshal(src)
if err != nil {
t.Fatalf("#%d: marshaling to JSON failed: %v", i, err)
}
if want, got := tt.Expected, string(data); want != got {
t.Fatalf("#%d: expected\n%s\ngot:\n%s", i, want, got)
}
}
}

0 comments on commit 987f2d7

Please sign in to comment.