From f46f6950ec611aefe220c1b21cdaef8b207a0764 Mon Sep 17 00:00:00 2001
From: Yevhen Tienkaiev <hronom@gmail.com>
Date: Thu, 29 Jul 2021 21:28:29 +0300
Subject: [PATCH] Rank feature - unknown field linear (#983)

Signed-off-by: Yevhen Tienkaiev <hronom@gmail.com>
---
 .../QueryDSLDocumentationTests.java           |  6 +++
 .../index/query/RankFeatureQueryBuilder.java  | 52 ++++++++++++++++++-
 .../index/query/RankFeatureQueryBuilders.java |  9 ++++
 .../query/RankFeatureQueryBuilderTests.java   |  5 +-
 4 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/client/rest-high-level/src/test/java/org/opensearch/client/documentation/QueryDSLDocumentationTests.java b/client/rest-high-level/src/test/java/org/opensearch/client/documentation/QueryDSLDocumentationTests.java
index 7b3d21cbaf5..135a0bea730 100644
--- a/client/rest-high-level/src/test/java/org/opensearch/client/documentation/QueryDSLDocumentationTests.java
+++ b/client/rest-high-level/src/test/java/org/opensearch/client/documentation/QueryDSLDocumentationTests.java
@@ -493,4 +493,10 @@ public class QueryDSLDocumentationTests extends OpenSearchTestCase {
             0.6f
         );
     }
+
+    public void testRankFeatureLinear() {
+        RankFeatureQueryBuilders.linear(
+            "pagerank"
+        );
+    }
 }
diff --git a/modules/mapper-extras/src/main/java/org/opensearch/index/query/RankFeatureQueryBuilder.java b/modules/mapper-extras/src/main/java/org/opensearch/index/query/RankFeatureQueryBuilder.java
index a2301954467..be38eeb0065 100644
--- a/modules/mapper-extras/src/main/java/org/opensearch/index/query/RankFeatureQueryBuilder.java
+++ b/modules/mapper-extras/src/main/java/org/opensearch/index/query/RankFeatureQueryBuilder.java
@@ -39,6 +39,7 @@ import org.opensearch.common.ParseField;
 import org.opensearch.common.io.stream.StreamInput;
 import org.opensearch.common.io.stream.StreamOutput;
 import org.opensearch.common.xcontent.ConstructingObjectParser;
+import org.opensearch.common.xcontent.ObjectParser;
 import org.opensearch.common.xcontent.XContentBuilder;
 import org.opensearch.index.mapper.RankFeatureFieldMapper.RankFeatureFieldType;
 import org.opensearch.index.mapper.RankFeatureMetaFieldMapper;
@@ -257,6 +258,51 @@ public final class RankFeatureQueryBuilder extends AbstractQueryBuilder<RankFeat
                 return FeatureField.newSigmoidQuery(field, feature, DEFAULT_BOOST, pivot, exp);
             }
         }
+
+        /**
+         * A scoring function that scores documents as simply {@code S}
+         * where S is the indexed value of the static feature.
+         */
+        public static class Linear extends ScoreFunction {
+
+            private static final ObjectParser<Linear, Void> PARSER = new ObjectParser<>("linear", Linear::new);
+
+            public Linear() {
+            }
+
+            private Linear(StreamInput in) {
+                this();
+            }
+
+            @Override
+            public boolean equals(Object obj) {
+                if (obj == null || getClass() != obj.getClass()) {
+                    return false;
+                }
+                return true;
+            }
+
+            @Override
+            public int hashCode() {
+                return getClass().hashCode();
+            }
+
+            @Override
+            void writeTo(StreamOutput out) throws IOException {
+                out.writeByte((byte) 3);
+            }
+
+            @Override
+            void doXContent(XContentBuilder builder) throws IOException {
+                builder.startObject("linear");
+                builder.endObject();
+            }
+
+            @Override
+            Query toQuery(String field, String feature, boolean positiveScoreImpact) throws IOException {
+                return FeatureField.newLinearQuery(field, feature, DEFAULT_BOOST);
+            }
+        }
     }
 
     private static ScoreFunction readScoreFunction(StreamInput in) throws IOException {
@@ -268,6 +314,8 @@ public final class RankFeatureQueryBuilder extends AbstractQueryBuilder<RankFeat
             return new ScoreFunction.Saturation(in);
         case 2:
             return new ScoreFunction.Sigmoid(in);
+        case 3:
+            return new ScoreFunction.Linear(in);
         default:
             throw new IOException("Illegal score function id: " + b);
         }
@@ -281,7 +329,7 @@ public final class RankFeatureQueryBuilder extends AbstractQueryBuilder<RankFeat
                 long numNonNulls = Arrays.stream(args, 3, args.length).filter(Objects::nonNull).count();
                 final RankFeatureQueryBuilder query;
                 if (numNonNulls > 1) {
-                    throw new IllegalArgumentException("Can only specify one of [log], [saturation] and [sigmoid]");
+                    throw new IllegalArgumentException("Can only specify one of [log], [saturation], [sigmoid] and [linear]");
                 } else if (numNonNulls == 0) {
                     query = new RankFeatureQueryBuilder(field, new ScoreFunction.Saturation());
                 } else {
@@ -305,6 +353,8 @@ public final class RankFeatureQueryBuilder extends AbstractQueryBuilder<RankFeat
                 ScoreFunction.Saturation.PARSER, new ParseField("saturation"));
         PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(),
                 ScoreFunction.Sigmoid.PARSER, new ParseField("sigmoid"));
+        PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(),
+            ScoreFunction.Linear.PARSER, new ParseField("linear"));
     }
 
     public static final String NAME = "rank_feature";
diff --git a/modules/mapper-extras/src/main/java/org/opensearch/index/query/RankFeatureQueryBuilders.java b/modules/mapper-extras/src/main/java/org/opensearch/index/query/RankFeatureQueryBuilders.java
index cc9d81c2bca..853bf0ff4e8 100644
--- a/modules/mapper-extras/src/main/java/org/opensearch/index/query/RankFeatureQueryBuilders.java
+++ b/modules/mapper-extras/src/main/java/org/opensearch/index/query/RankFeatureQueryBuilders.java
@@ -77,4 +77,13 @@ public final class RankFeatureQueryBuilders {
         return new RankFeatureQueryBuilder(fieldName, new RankFeatureQueryBuilder.ScoreFunction.Sigmoid(pivot, exp));
     }
 
+    /**
+     * Return a new {@link RankFeatureQueryBuilder} that will score documents as
+     * {@code S)} where S is the indexed value of the static feature.
+     * @param fieldName     field that stores features
+     */
+    public static RankFeatureQueryBuilder linear(String fieldName) {
+        return new RankFeatureQueryBuilder(fieldName, new RankFeatureQueryBuilder.ScoreFunction.Linear());
+    }
+
 }
diff --git a/modules/mapper-extras/src/test/java/org/opensearch/index/query/RankFeatureQueryBuilderTests.java b/modules/mapper-extras/src/test/java/org/opensearch/index/query/RankFeatureQueryBuilderTests.java
index ca3b083f54d..a25219b425d 100644
--- a/modules/mapper-extras/src/test/java/org/opensearch/index/query/RankFeatureQueryBuilderTests.java
+++ b/modules/mapper-extras/src/test/java/org/opensearch/index/query/RankFeatureQueryBuilderTests.java
@@ -73,7 +73,7 @@ public class RankFeatureQueryBuilderTests extends AbstractQueryTestCase<RankFeat
     protected RankFeatureQueryBuilder doCreateTestQueryBuilder() {
         ScoreFunction function;
         boolean mayUseNegativeField = true;
-        switch (random().nextInt(3)) {
+        switch (random().nextInt(4)) {
         case 0:
             mayUseNegativeField = false;
             function = new ScoreFunction.Log(1 + randomFloat());
@@ -88,6 +88,9 @@ public class RankFeatureQueryBuilderTests extends AbstractQueryTestCase<RankFeat
         case 2:
             function = new ScoreFunction.Sigmoid(randomFloat(), randomFloat());
             break;
+        case 3:
+            function = new ScoreFunction.Linear();
+            break;
         default:
             throw new AssertionError();
         }