Support semantic reranking using contextual snippets instead of entire field text

docs/changelog/129369.yaml

@@ -0,0 +1,6 @@
+pr: 129369
+summary: Support semantic reranking using contextual snippets instead of entire field
+  text
+area: Relevance
+type: enhancement
+issues: []

qa/ccs-common-rest/src/yamlRestTest/java/org/elasticsearch/test/rest/yaml/CcsCommonYamlTestSuiteIT.java

@@ -92,7 +92,8 @@ public class CcsCommonYamlTestSuiteIT extends ESClientYamlSuiteTestCase {
         .feature(FeatureFlag.TIME_SERIES_MODE)
         .feature(FeatureFlag.SUB_OBJECTS_AUTO_ENABLED)
         .feature(FeatureFlag.IVF_FORMAT)
-        .feature(FeatureFlag.SYNTHETIC_VECTORS);
+        .feature(FeatureFlag.SYNTHETIC_VECTORS)
+        .feature(FeatureFlag.RERANK_SNIPPETS);
 
     private static ElasticsearchCluster remoteCluster = ElasticsearchCluster.local()
         .name(REMOTE_CLUSTER_NAME)

qa/ccs-common-rest/src/yamlRestTest/java/org/elasticsearch/test/rest/yaml/RcsCcsCommonYamlTestSuiteIT.java

@@ -94,6 +94,7 @@ public class RcsCcsCommonYamlTestSuiteIT extends ESClientYamlSuiteTestCase {
         .feature(FeatureFlag.SUB_OBJECTS_AUTO_ENABLED)
         .feature(FeatureFlag.IVF_FORMAT)
         .feature(FeatureFlag.SYNTHETIC_VECTORS)
+        .feature(FeatureFlag.RERANK_SNIPPETS)
         .user("test_admin", "x-pack-test-password");
 
     private static ElasticsearchCluster fulfillingCluster = ElasticsearchCluster.local()

qa/smoke-test-multinode/src/yamlRestTest/java/org/elasticsearch/smoketest/SmokeTestMultiNodeClientYamlTestSuiteIT.java

@@ -40,6 +40,7 @@ public class SmokeTestMultiNodeClientYamlTestSuiteIT extends ESClientYamlSuiteTe
         .feature(FeatureFlag.USE_LUCENE101_POSTINGS_FORMAT)
         .feature(FeatureFlag.IVF_FORMAT)
         .feature(FeatureFlag.SYNTHETIC_VECTORS)
+        .feature(FeatureFlag.RERANK_SNIPPETS)
         .build();
 
     public SmokeTestMultiNodeClientYamlTestSuiteIT(@Name("yaml") ClientYamlTestCandidate testCandidate) {

rest-api-spec/src/yamlRestTest/java/org/elasticsearch/test/rest/ClientYamlTestSuiteIT.java

@@ -40,6 +40,7 @@ public class ClientYamlTestSuiteIT extends ESClientYamlSuiteTestCase {
         .feature(FeatureFlag.USE_LUCENE101_POSTINGS_FORMAT)
         .feature(FeatureFlag.IVF_FORMAT)
         .feature(FeatureFlag.SYNTHETIC_VECTORS)
+        .feature(FeatureFlag.RERANK_SNIPPETS)
         .build();
 
     public ClientYamlTestSuiteIT(@Name("yaml") ClientYamlTestCandidate testCandidate) {

server/src/internalClusterTest/java/org/elasticsearch/search/rank/FieldBasedRerankerIT.java

@@ -193,7 +193,7 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
                         RankFeatureDoc[] rankFeatureDocs = new RankFeatureDoc[hits.getHits().length];
                         for (int i = 0; i < hits.getHits().length; i++) {
                             rankFeatureDocs[i] = new RankFeatureDoc(hits.getHits()[i].docId(), hits.getHits()[i].getScore(), shardId);
-                            rankFeatureDocs[i].featureData(hits.getHits()[i].field(field).getValue().toString());
+                            rankFeatureDocs[i].featureData(List.of(hits.getHits()[i].field(field).getValue().toString()));
                         }
                         return new RankFeatureShardResult(rankFeatureDocs);
                     } catch (Exception ex) {
@@ -210,7 +210,7 @@ public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorCo
                 protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
                     float[] scores = new float[featureDocs.length];
                     for (int i = 0; i < featureDocs.length; i++) {
-                        scores[i] = Float.parseFloat(featureDocs[i].featureData);
+                        scores[i] = Float.parseFloat(featureDocs[i].featureData.get(0));
                     }
                     scoreListener.onResponse(scores);
                 }

server/src/internalClusterTest/java/org/elasticsearch/search/rank/MockedRequestActionBasedRerankerIT.java

@@ -275,7 +275,7 @@ protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[
                 l.onResponse(scores);
             });
 
-            List<String> featureData = Arrays.stream(featureDocs).map(x -> x.featureData).toList();
+            List<String> featureData = Arrays.stream(featureDocs).map(x -> x.featureData).flatMap(List::stream).toList();
             TestRerankingActionRequest request = generateRequest(featureData);
             try {
                 ActionType<TestRerankingActionResponse> action = actionType();

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

@@ -348,6 +348,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion SHARD_WRITE_LOAD_IN_CLUSTER_INFO = def(9_126_0_00);
     public static final TransportVersion ESQL_SAMPLE_OPERATOR_STATUS = def(9_127_0_00);
     public static final TransportVersion ESQL_TOPN_TIMINGS = def(9_128_0_00);
+    public static final TransportVersion RERANK_SNIPPETS = def(9_129_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -1187,9 +1187,11 @@ public SearchSourceBuilder rewrite(QueryRewriteContext context) throws IOExcepti
                 sliceBuilder,
                 sorts,
                 rescoreBuilders,
-                highlightBuilder
+                highlightBuilder,
+                rankBuilder
             )
         ));
+
         if (retrieverBuilder != null) {
             var newRetriever = retrieverBuilder.rewrite(context);
             if (newRetriever != retrieverBuilder) {
@@ -1205,6 +1207,11 @@ public SearchSourceBuilder rewrite(QueryRewriteContext context) throws IOExcepti
             }
         }
 
+        RankBuilder rankBuilder = null;
+        if (this.rankBuilder != null) {
+            rankBuilder = this.rankBuilder.rewrite(context);
+        }
+
         List<SubSearchSourceBuilder> subSearchSourceBuilders = Rewriteable.rewrite(this.subSearchSourceBuilders, context);
         QueryBuilder postQueryBuilder = null;
         if (this.postQueryBuilder != null) {
@@ -1229,7 +1236,8 @@ public SearchSourceBuilder rewrite(QueryRewriteContext context) throws IOExcepti
             || aggregations != this.aggregations
             || rescoreBuilders != this.rescoreBuilders
             || sorts != this.sorts
-            || this.highlightBuilder != highlightBuilder;
+            || this.highlightBuilder != highlightBuilder
+            || this.rankBuilder != rankBuilder;
         if (rewritten) {
             return shallowCopy(
                 subSearchSourceBuilders,
@@ -1239,7 +1247,8 @@ public SearchSourceBuilder rewrite(QueryRewriteContext context) throws IOExcepti
                 this.sliceBuilder,
                 sorts,
                 rescoreBuilders,
-                highlightBuilder
+                highlightBuilder,
+                rankBuilder
             );
         }
         return this;
@@ -1257,7 +1266,8 @@ public SearchSourceBuilder shallowCopy() {
             sliceBuilder,
             sorts,
             rescoreBuilders,
-            highlightBuilder
+            highlightBuilder,
+            rankBuilder
         );
     }
 
@@ -1274,7 +1284,8 @@ private SearchSourceBuilder shallowCopy(
         SliceBuilder slice,
         List<SortBuilder<?>> sorts,
         List<RescorerBuilder> rescoreBuilders,
-        HighlightBuilder highlightBuilder
+        HighlightBuilder highlightBuilder,
+        RankBuilder rankBuilder
     ) {
         SearchSourceBuilder rewrittenBuilder = new SearchSourceBuilder();
         rewrittenBuilder.aggregations = aggregations;

server/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/HighlightBuilder.java

@@ -399,7 +399,7 @@ public Field(String name) {
             this.name = name;
         }
 
-        private Field(Field template, QueryBuilder builder) {
+        Field(Field template, QueryBuilder builder) {
             super(template, builder);
             name = template.name;
             fragmentOffset = template.fragmentOffset;

server/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/SearchHighlightContext.java

@@ -40,7 +40,7 @@ public static class Field {
         private final String field;
         private final FieldOptions fieldOptions;
 
-        Field(String field, FieldOptions fieldOptions) {
+        public Field(String field, FieldOptions fieldOptions) {
             assert field != null;
             assert fieldOptions != null;
             this.field = field;

server/src/main/java/org/elasticsearch/search/rank/RankBuilder.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.UpdateForV10;
 import org.elasticsearch.features.NodeFeature;
+import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.rank.context.QueryPhaseRankCoordinatorContext;
@@ -60,6 +61,10 @@ public final void writeTo(StreamOutput out) throws IOException {
         doWriteTo(out);
     }
 
+    public RankBuilder rewrite(QueryRewriteContext queryRewriteContext) throws IOException {
+        return this;
+    }
+
     protected abstract void doWriteTo(StreamOutput out) throws IOException;
 
     @Override

server/src/main/java/org/elasticsearch/search/rank/context/RankFeaturePhaseRankShardContext.java

@@ -11,6 +11,7 @@
 
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.search.SearchHits;
+import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.rank.RankShardResult;
 
 /**
@@ -37,4 +38,13 @@ public String getField() {
      */
     @Nullable
     public abstract RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId);
+
+    /**
+     * Prepares a SearchContext with any additional information needed before executing
+     * commands on shards.
+     * @param context SearchContext
+     */
+    public void prepareForFetch(SearchContext context) {
+        // Default no-op
+    }
 }

server/src/main/java/org/elasticsearch/search/rank/feature/RankFeatureDoc.java

@@ -10,12 +10,14 @@
 package org.elasticsearch.search.rank.feature;
 
 import org.apache.lucene.search.Explanation;
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.search.rank.RankDoc;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Objects;
 
 /**
@@ -26,29 +28,38 @@ public class RankFeatureDoc extends RankDoc {
     public static final String NAME = "rank_feature_doc";
 
     // TODO: update to support more than 1 fields; and not restrict to string data
-    public String featureData;
+    public List<String> featureData;
 
     public RankFeatureDoc(int doc, float score, int shardIndex) {
         super(doc, score, shardIndex);
     }
 
     public RankFeatureDoc(StreamInput in) throws IOException {
         super(in);
-        featureData = in.readOptionalString();
+        if (in.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
+            featureData = in.readOptionalStringCollectionAsList();
+        } else {
+            String featureDataString = in.readOptionalString();
+            featureData = featureDataString == null ? null : List.of(featureDataString);
+        }
     }
 
     @Override
     public Explanation explain(Explanation[] sources, String[] queryNames) {
         throw new UnsupportedOperationException("explain is not supported for {" + getClass() + "}");
     }
 
-    public void featureData(String featureData) {
+    public void featureData(List<String> featureData) {
         this.featureData = featureData;
     }
 
     @Override
     protected void doWriteTo(StreamOutput out) throws IOException {
-        out.writeOptionalString(featureData);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
+            out.writeOptionalStringCollection(featureData);
+        } else {
+            out.writeOptionalString(featureData.get(0));
+        }
     }
 
     @Override
@@ -59,7 +70,7 @@ protected boolean doEquals(RankDoc rd) {
 
     @Override
     protected int doHashCode() {
-        return Objects.hashCode(featureData);
+        return Objects.hash(featureData);
     }
 
     @Override
@@ -69,6 +80,6 @@ public String getWriteableName() {
 
     @Override
     protected void doToXContent(XContentBuilder builder, Params params) throws IOException {
-        builder.field("featureData", featureData);
+        builder.array("featureData", featureData);
     }
 }

server/src/main/java/org/elasticsearch/search/rank/feature/RankFeatureShardPhase.java

@@ -48,10 +48,10 @@ public static void prepareForFetch(SearchContext searchContext, RankFeatureShard
 
         RankFeaturePhaseRankShardContext rankFeaturePhaseRankShardContext = shardContext(searchContext);
         if (rankFeaturePhaseRankShardContext != null) {
-            assert rankFeaturePhaseRankShardContext.getField() != null : "field must not be null";
-            searchContext.fetchFieldsContext(
-                new FetchFieldsContext(Collections.singletonList(new FieldAndFormat(rankFeaturePhaseRankShardContext.getField(), null)))
-            );
+            String field = rankFeaturePhaseRankShardContext.getField();
+            assert field != null : "field must not be null";
+            searchContext.fetchFieldsContext(new FetchFieldsContext(Collections.singletonList(new FieldAndFormat(field, null))));
+            rankFeaturePhaseRankShardContext.prepareForFetch(searchContext);
             searchContext.storedFieldsContext(StoredFieldsContext.fromList(Collections.singletonList(StoredFieldsContext._NONE_)));
             searchContext.addFetchResult();
             Arrays.sort(request.getDocIds());

server/src/main/java/org/elasticsearch/search/rank/feature/RerankSnippetInput.java

@@ -0,0 +1,90 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.search.rank.feature;
+
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.index.query.QueryBuilder;
+
+import java.io.IOException;
+import java.util.Objects;
+
+public class RerankSnippetInput implements Writeable {
+
+    public final Integer numSnippets;
+    private final String inferenceText;
+    private final Integer tokenSizeLimit;
+    public final QueryBuilder snippetQueryBuilder;
+
+    public static final int DEFAULT_NUM_SNIPPETS = 1;
+
+    public RerankSnippetInput(StreamInput in) throws IOException {
+        this.numSnippets = in.readOptionalVInt();
+        this.inferenceText = in.readString();
+        this.tokenSizeLimit = in.readOptionalVInt();
+        this.snippetQueryBuilder = in.readOptionalNamedWriteable(QueryBuilder.class);
+    }
+
+    public RerankSnippetInput(Integer numSnippets) {
+        this(numSnippets, null, null);
+    }
+
+    public RerankSnippetInput(Integer numSnippets, String inferenceText, Integer tokenSizeLimit) {
+        this(numSnippets, inferenceText, tokenSizeLimit, null);
+    }
+
+    public RerankSnippetInput(Integer numSnippets, String inferenceText, Integer tokenSizeLimit, QueryBuilder snippetQueryBuilder) {
+        this.numSnippets = numSnippets;
+        this.inferenceText = inferenceText;
+        this.tokenSizeLimit = tokenSizeLimit;
+        this.snippetQueryBuilder = snippetQueryBuilder;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeOptionalVInt(numSnippets);
+        out.writeString(inferenceText);
+        out.writeOptionalVInt(tokenSizeLimit);
+        out.writeOptionalNamedWriteable(snippetQueryBuilder);
+    }
+
+    public Integer numSnippets() {
+        return numSnippets;
+    }
+
+    public String inferenceText() {
+        return inferenceText;
+    }
+
+    public Integer tokenSizeLimit() {
+        return tokenSizeLimit;
+    }
+
+    public QueryBuilder snippetQueryBuilder() {
+        return snippetQueryBuilder;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        RerankSnippetInput that = (RerankSnippetInput) o;
+        return Objects.equals(numSnippets, that.numSnippets)
+            && Objects.equals(inferenceText, that.inferenceText)
+            && Objects.equals(tokenSizeLimit, that.tokenSizeLimit)
+            && Objects.equals(snippetQueryBuilder, that.snippetQueryBuilder);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(numSnippets, inferenceText, tokenSizeLimit, snippetQueryBuilder);
+    }
+}