New rescorer based on script

docs/changelog/74274.yaml

@@ -0,0 +1,6 @@
+pr: 74274
+summary: Introduce new rescorer based on script
+area: Search
+type: feature
+issues:
+  - 52338

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search/620_rescore_script.yml

@@ -0,0 +1,277 @@
+setup:
+  - requires:
+      cluster_features: [ "search.rescore.script" ]
+      reason: "Test requires feature 'search.rescore.script' to test script rescore"
+      test_runner_features: "close_to"
+  - skip:
+      reason: "contains is a newly added assertion"
+      features: contains
+
+  - do:
+      indices.create:
+        index: books
+        body:
+          settings:
+            number_of_shards: 1
+          mappings:
+            properties:
+              title:
+                type: text
+              author:
+                type: keyword
+              num_likes:
+                type: integer
+
+  - do:
+      bulk:
+        index: books
+        refresh: true
+        body:
+          - '{"index": {"_id": "1"}}'
+          - '{"title": "The Ethics of Ambiguity", "author": "Simone de Beauvoir", "num_likes": 150}'
+          - '{"index": {"_id": "2"}}'
+          - '{"title": "The Ethics of Being and Nothingness", "author": "Jean-Paul Sartre", "num_likes": 250}'
+          - '{"index": {"_id": "3"}}'
+          - '{"title": "The Ambiguity of The Second Sex", "author": "Simone de Beauvoir", "num_likes": 300}'
+          - '{"index": {"_id": "4"}}'
+          - '{"title": "The Ethics of Nausea", "author": "Jean-Paul Sartre", "num_likes": 180}'
+          - '{"index": {"_id": "5"}}'
+          - '{"title": "The Ambiguity in the Myth of Sisyphus", "author": "Albert Camus", "num_likes": 220}'
+
+---
+"Basic script rescore":
+
+  # simple rescore script for all top docs
+  - do:
+      search:
+        index: books
+        body:
+          query:
+            match:
+             title: "ethics of ambiguity"
+          rescore:
+            window_size: 5
+            script:
+              script:
+                source: "doc['num_likes'].value * params.multiplier"
+                params:
+                  multiplier: 10
+  - match: { hits.total.value: 5 }
+  - match: { hits.hits.0._id: "3" }
+  - close_to: { hits.hits.0._score:  {value: 3000, error: 0.0001} }
+  - match: { hits.hits.1._id: "2" }
+  - close_to: { hits.hits.1._score:  {value: 2500, error: 0.0001} }
+  - match: { hits.hits.2._id: "5" }
+  - close_to: { hits.hits.2._score:  {value: 2200, error: 0.0001} }
+  - match: { hits.hits.3._id: "4" }
+  - close_to: { hits.hits.3._score:  {value: 1800, error: 0.0001} }
+  - match: { hits.hits.4._id: "1" }
+  - close_to: { hits.hits.4._score:  {value: 1500, error: 0.0001} }
+
+  # rescore script with _score for all top docs
+  - do:
+      search:
+        index: books
+        body:
+          query:
+            match:
+              title: "ethics of ambiguity"
+          rescore:
+            window_size: 5
+            script:
+              script:
+                source: "doc['num_likes'].value * params.multiplier + _score"
+                params:
+                  multiplier: 10
+  - match: { hits.hits.0._id: "3" }
+  - close_to: { hits.hits.0._score: { value: 3000.5989, error: 0.0001 } }
+  - match: { hits.hits.1._id: "2" }
+  - close_to: { hits.hits.1._score: { value: 2500.5989, error: 0.0001 } }
+  - match: { hits.hits.2._id: "5" }
+  - close_to: { hits.hits.2._score: { value: 2200.5583, error: 0.0001 } }
+  - match: { hits.hits.3._id: "4" }
+  - close_to: { hits.hits.3._score: { value: 1800.7003, error: 0.0001 } }
+  - match: { hits.hits.4._id: "1" }
+  - close_to: { hits.hits.4._score: { value: 1501.3032, error: 0.0001 } }
+
+  # simple rescore script for small window size
+  - do:
+      search:
+        index: books
+        body:
+          query:
+            match:
+              title: "ethics of ambiguity"
+          rescore:
+            window_size: 3
+            script:
+              script:
+                source: "doc['num_likes'].value * params.multiplier"
+                params:
+                  multiplier: 10
+  - match: { hits.hits.0._id: "2" }
+  - close_to: { hits.hits.0._score: { value: 2500, error: 0.0001 } }
+  - match: { hits.hits.1._id: "4" }
+  - close_to: { hits.hits.1._score: { value: 1800, error: 0.0001 } }
+  - match: { hits.hits.2._id: "1" }
+  - close_to: { hits.hits.2._score: { value: 1500, error: 0.0001 } }
+  - match: { hits.hits.3._id: "3" }
+  - close_to: { hits.hits.3._score: { value: 0.59879, error: 0.0001 } }
+  - match: { hits.hits.4._id: "5" }
+  - close_to: { hits.hits.4._score: { value: 0.5583, error: 0.0001 } }
+
+  # rescore script with _score for small window size
+  - do:
+      search:
+        index: books
+        body:
+          query:
+            match:
+              title: "ethics of ambiguity"
+          rescore:
+            window_size: 3
+            script:
+              script:
+                source: "doc['num_likes'].value * params.multiplier + _score"
+                params:
+                  multiplier: 10
+  - match: { hits.hits.0._id: "2" }
+  - close_to: { hits.hits.0._score: { value: 2500.5989, error: 0.0001 } }
+  - match: { hits.hits.1._id: "4" }
+  - close_to: { hits.hits.1._score: { value: 1800.7003, error: 0.0001 } }
+  - match: { hits.hits.2._id: "1" }
+  - close_to: { hits.hits.2._score: { value: 1501.3032, error: 0.0001 } }
+  - match: { hits.hits.3._id: "3" }
+  - close_to: { hits.hits.3._score: { value: 0.59879, error: 0.0001 } }
+  - match: { hits.hits.4._id: "5" }
+  - close_to: { hits.hits.4._score: { value: 0.5583, error: 0.0001 } }
+
+---
+"Multiple script rescore":
+  - do:
+      search:
+        index: books
+        body:
+          query:
+            match:
+              title: "ethics of ambiguity"
+          rescore:
+            - window_size: 3
+              script:
+                script:
+                  source: "doc['num_likes'].value * params.multiplier + _score"
+                  params:
+                    multiplier: 10
+            - window_size: 2
+              script:
+                script:
+                  source: "_score * params.factor"
+                  params:
+                    factor: 10
+  - match: { hits.hits.0._id: "2" }
+  - close_to: { hits.hits.0._score: { value: 25005.989, error: 0.005 } } # rescored by 1st and 2nd scripts
+  - match: { hits.hits.1._id: "4" }
+  - close_to: { hits.hits.1._score: { value: 18007.0039, error: 0.005 } } # rescored by 1st and 2nd scripts
+  - match: { hits.hits.2._id: "1" }
+  - close_to: { hits.hits.2._score: { value: 1501.3032, error: 0.0005 } } # rescored by 1st script
+  - match: { hits.hits.3._id: "3" }
+  - close_to: { hits.hits.3._score: { value: 0.59879, error: 0.0001 } } # not rescored
+  - match: { hits.hits.4._id: "5" }
+  - close_to: { hits.hits.4._score: { value: 0.5583, error: 0.0001 } } # not rescored
+
+---
+"Rescore Script With Explanation":
+  - do:
+      search:
+        index: books
+        explain: true
+        body:
+          query:
+            match:
+              title: "ethics of ambiguity"
+          rescore:
+            window_size: 3
+            script:
+              script:
+                source: "doc['num_likes'].value * params.multiplier + _score"
+                params:
+                  multiplier: 10
+  - match: { hits.total.value: 5 }
+
+  # hit went through rescore
+  - contains: { hits.hits.0._explanation.description: "script score function" }
+  - close_to: { hits.hits.0._explanation.value: { value: 2500.5989, error: 0.0001 } }
+  - match: { hits.hits.0._explanation.details.0.description: '_score: ' }
+  - close_to: { hits.hits.0._explanation.details.0.value: { value: 0.5987902, error: 0.0001 } }
+
+  # hit did not go through rescore
+  - match: { hits.hits.4._explanation.description: "sum of:" }
+  - close_to: { hits.hits.4._explanation.value: { value: 0.5583, error: 0.0001 } }
+
+
+---
+"Script rescore on Multiple Segments":
+  # update some documents to create multiple segments
+  - do:
+      bulk:
+        index: books
+        refresh: true
+        body:
+          - '{"update": {"_id": "1"}}'
+          - '{"doc": {"num_likes": 300}}'
+          - '{"update": {"_id": "3"}}'
+          - '{"doc": {"num_likes": 150}}'
+
+  - do:
+      search:
+        index: books
+        body:
+          query:
+            match:
+              title: "ethics of ambiguity"
+          rescore:
+            window_size: 5
+            script:
+              script:
+                source: "doc['num_likes'].value * params.multiplier + _score"
+                params:
+                  multiplier: 10
+  - match: { hits.total.value: 5 }
+  - match: { hits.hits.0._id: "1" }
+  - close_to: { hits.hits.0._score: { value: 3001.1267, error: 0.0001 } }
+  - match: { hits.hits.1._id: "2" }
+  - close_to: { hits.hits.1._score: { value: 2500.6064, error: 0.0001 } }
+  - match: { hits.hits.2._id: "5" }
+  - close_to: { hits.hits.2._score: { value: 2200.3877, error: 0.0001 } }
+  - match: { hits.hits.3._id: "4" }
+  - close_to: { hits.hits.3._score: { value: 1800.7106, error: 0.0001 } }
+  - match: { hits.hits.4._id: "3" }
+  - close_to: { hits.hits.4._score: { value: 1500.4163, error: 0.0001 } }
+
+  # rescore script with _score for small window size
+  - do:
+      search:
+        index: books
+        body:
+          query:
+            match:
+              title: "ethics of ambiguity"
+          rescore:
+            window_size: 3
+            script:
+              script:
+                source: "doc['num_likes'].value * params.multiplier + _score"
+                params:
+                  multiplier: 10
+  - match: { hits.total.value: 5 }
+  - match: { hits.hits.0._id: "1" }
+  - close_to: { hits.hits.0._score: { value: 3001.1267, error: 0.0001 } }
+  - match: { hits.hits.1._id: "2" }
+  - close_to: { hits.hits.1._score: { value: 2500.6064, error: 0.0001 } }
+  - match: { hits.hits.2._id: "4" }
+  - close_to: { hits.hits.2._score: { value: 1800.7106, error: 0.0001 } }
+  - match: { hits.hits.3._id: "3" }
+  - close_to: { hits.hits.3._score: { value: 0.41622, error: 0.0001 } }
+  - match: { hits.hits.4._id: "5" }
+  - close_to: { hits.hits.4._score: { value: 0.38778, error: 0.0001 } }
+

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -366,6 +366,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion SIMULATE_INGEST_EFFECTIVE_MAPPING = def(9_140_0_00);
     public static final TransportVersion RESOLVE_INDEX_MODE_ADDED = def(9_141_0_00);
     public static final TransportVersion DATA_STREAM_WRITE_INDEX_ONLY_SETTINGS = def(9_142_0_00);
+    public static final TransportVersion SCRIPT_RESCORER = def(9_143_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/common/lucene/search/function/ScriptScoreQuery.java

@@ -386,4 +386,12 @@ public long cost() {
 
     }
 
+    public boolean needsScore() {
+        return scriptBuilder.needs_score();
+    }
+
+    public ScriptScoreQuery cloneWithNewSubQuery(Query newSubQuery) {
+        return new ScriptScoreQuery(newSubQuery, script, scriptBuilder, lookup, minScore, indexName, shardId, indexVersion);
+    }
+
 }

server/src/main/java/org/elasticsearch/search/SearchFeatures.java

@@ -33,6 +33,7 @@ public Set<NodeFeature> getFeatures() {
     static final NodeFeature MULTI_MATCH_CHECKS_POSITIONS = new NodeFeature("search.multi.match.checks.positions");
     public static final NodeFeature BBQ_HNSW_DEFAULT_INDEXING = new NodeFeature("search.vectors.mappers.default_bbq_hnsw");
     public static final NodeFeature SEARCH_WITH_NO_DIMENSIONS_BUGFIX = new NodeFeature("search.vectors.no_dimensions_bugfix");
+    public static final NodeFeature SEARCH_RESCORE_SCRIPT = new NodeFeature("search.rescore.script");
 
     @Override
     public Set<NodeFeature> getTestFeatures() {
@@ -43,7 +44,8 @@ public Set<NodeFeature> getTestFeatures() {
             INT_SORT_FOR_INT_SHORT_BYTE_FIELDS,
             MULTI_MATCH_CHECKS_POSITIONS,
             BBQ_HNSW_DEFAULT_INDEXING,
-            SEARCH_WITH_NO_DIMENSIONS_BUGFIX
+            SEARCH_WITH_NO_DIMENSIONS_BUGFIX,
+            SEARCH_RESCORE_SCRIPT
         );
     }
 }

server/src/main/java/org/elasticsearch/search/SearchModule.java

@@ -232,6 +232,7 @@
 import org.elasticsearch.search.rank.feature.RankFeatureShardResult;
 import org.elasticsearch.search.rescore.QueryRescorerBuilder;
 import org.elasticsearch.search.rescore.RescorerBuilder;
+import org.elasticsearch.search.rescore.ScriptRescorerBuilder;
 import org.elasticsearch.search.retriever.KnnRetrieverBuilder;
 import org.elasticsearch.search.retriever.RescorerRetrieverBuilder;
 import org.elasticsearch.search.retriever.RetrieverBuilder;
@@ -825,6 +826,7 @@ private void registerPipelineAggregation(PipelineAggregationSpec spec) {
 
     private void registerRescorers(List<SearchPlugin> plugins) {
         registerRescorer(new RescorerSpec<>(QueryRescorerBuilder.NAME, QueryRescorerBuilder::new, QueryRescorerBuilder::fromXContent));
+        registerRescorer(new RescorerSpec<>(ScriptRescorerBuilder.NAME, ScriptRescorerBuilder::new, ScriptRescorerBuilder::fromXContent));
         registerFromPlugin(plugins, SearchPlugin::getRescorers, this::registerRescorer);
     }
 

server/src/main/java/org/elasticsearch/search/rescore/QueryRescorer.java

@@ -11,14 +11,12 @@
 
 import org.apache.lucene.search.Explanation;
 import org.apache.lucene.search.IndexSearcher;
-import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.TopDocs;
 import org.elasticsearch.index.query.ParsedQuery;
 
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.Comparator;
 import java.util.List;
 import java.util.Set;
 
@@ -61,7 +59,7 @@ protected float combine(float firstPassScore, boolean secondPassMatches, float s
         };
 
         // First take top slice of incoming docs, to be rescored:
-        TopDocs topNFirstPass = topN(topDocs, rescoreContext.getWindowSize());
+        TopDocs topNFirstPass = Rescorer.topN(topDocs, rescoreContext.getWindowSize());
 
         // Save doc IDs for which rescoring was applied to be used in score explanation
         Set<Integer> topNDocIDs = Arrays.stream(topNFirstPass.scoreDocs).map(scoreDoc -> scoreDoc.doc).collect(toUnmodifiableSet());
@@ -118,24 +116,6 @@ public Explanation explain(int topLevelDocId, IndexSearcher searcher, RescoreCon
         return prim;
     }
 
-    private static final Comparator<ScoreDoc> SCORE_DOC_COMPARATOR = (o1, o2) -> {
-        int cmp = Float.compare(o2.score, o1.score);
-        return cmp == 0 ? Integer.compare(o1.doc, o2.doc) : cmp;
-    };
-
-    /** Returns a new {@link TopDocs} with the topN from the incoming one, or the same TopDocs if the number of hits is already &lt;=
-     *  topN. */
-    private static TopDocs topN(TopDocs in, int topN) {
-        if (in.scoreDocs.length < topN) {
-            return in;
-        }
-
-        ScoreDoc[] subset = new ScoreDoc[topN];
-        System.arraycopy(in.scoreDocs, 0, subset, 0, topN);
-
-        return new TopDocs(in.totalHits, subset);
-    }
-
     /** Modifies incoming TopDocs (in) by replacing the top hits with resorted's hits, and then resorting all hits. */
     private static TopDocs combine(TopDocs in, TopDocs resorted, QueryRescoreContext ctx) {