Update API - remove max size, default num snippets, rename

Author: kderusso

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

@@ -63,12 +63,10 @@ public static void prepareForFetch(SearchContext searchContext, RankFeatureShard
                     HighlightBuilder highlightBuilder = new HighlightBuilder().field(field).preTags("").postTags("");
                     // Force sorting by score to ensure that the first snippet is always the highest score
                     highlightBuilder.order(HighlightBuilder.Order.SCORE);
-                    if (snippets.numFragments() != null) {
-                        highlightBuilder.numOfFragments(snippets.numFragments());
-                    }
-                    if (snippets.maxSize() != null) {
-                        highlightBuilder.fragmentSize(snippets.maxSize());
+                    if (snippets.numSnippets() != null) {
+                        highlightBuilder.numOfFragments(snippets.numSnippets());
                     }
+                    highlightBuilder.fragmentSize(20); // TODO use model limit
                     SearchHighlightContext searchHighlightContext = highlightBuilder.build(searchContext.getSearchExecutionContext());
                     searchContext.highlight(searchHighlightContext);
                 }

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

@@ -15,15 +15,22 @@
 
 import java.io.IOException;
 
-public record RerankSnippetInput(Integer numFragments, Integer maxSize) implements Writeable {
+public record RerankSnippetInput(Integer numSnippets) implements Writeable {
+
+    private static final int DEFAULT_NUM_SNIPPETS = 1;
+
+    public RerankSnippetInput {
+        if (numSnippets == null) {
+            numSnippets = DEFAULT_NUM_SNIPPETS;
+        }
+    }
 
     public RerankSnippetInput(StreamInput in) throws IOException {
-        this(in.readOptionalVInt(), in.readOptionalVInt());
+        this(in.readOptionalVInt());
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeOptionalVInt(numFragments);
-        out.writeOptionalVInt(maxSize);
+        out.writeOptionalVInt(numSnippets);
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankFeaturePhaseRankCoordinatorContext.java

@@ -69,7 +69,7 @@ protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[
 
             List<RankedDocsResults.RankedDoc> rankedDocs = ((RankedDocsResults) results).getRankedDocs();
             final float[] scores;
-            if (featureDocs.length > 0 && featureDocs[0].featureData != null && featureDocs[0].featureData.size() > 1) {
+            if (this.snippets != null) {
                 scores = extractScoresFromRankedSnippets(rankedDocs, featureDocs);
             } else {
                 scores = extractScoresFromRankedDocs(rankedDocs);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java

@@ -49,8 +49,7 @@ public class TextSimilarityRankRetrieverBuilder extends CompoundRetrieverBuilder
     public static final ParseField FIELD_FIELD = new ParseField("field");
     public static final ParseField FAILURES_ALLOWED_FIELD = new ParseField("allow_rerank_failures");
     public static final ParseField SNIPPETS_FIELD = new ParseField("snippets");
-    public static final ParseField NUM_FRAGMENTS_FIELD = new ParseField("num_fragments");
-    public static final ParseField MAX_SIZE_FIELD = new ParseField("max_size");
+    public static final ParseField NUM_SNIPPETS_FIELD = new ParseField("num_snippets");
 
     public static final ConstructingObjectParser<TextSimilarityRankRetrieverBuilder, RetrieverParserContext> PARSER =
         new ConstructingObjectParser<>(TextSimilarityRankBuilder.NAME, args -> {
@@ -75,9 +74,8 @@ public class TextSimilarityRankRetrieverBuilder extends CompoundRetrieverBuilder
 
     private static final ConstructingObjectParser<RerankSnippetInput, RetrieverParserContext> SNIPPETS_PARSER =
         new ConstructingObjectParser<>(SNIPPETS_FIELD.getPreferredName(), true, args -> {
-            Integer numFragments = (Integer) args[0];
-            Integer maxSize = (Integer) args[1];
-            return new RerankSnippetInput(numFragments, maxSize);
+            Integer numSnippets = (Integer) args[0];
+            return new RerankSnippetInput(numSnippets);
         });
 
     static {
@@ -92,8 +90,7 @@ public class TextSimilarityRankRetrieverBuilder extends CompoundRetrieverBuilder
         PARSER.declareInt(optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD);
         PARSER.declareBoolean(optionalConstructorArg(), FAILURES_ALLOWED_FIELD);
         PARSER.declareObject(optionalConstructorArg(), SNIPPETS_PARSER, SNIPPETS_FIELD);
-        SNIPPETS_PARSER.declareInt(optionalConstructorArg(), NUM_FRAGMENTS_FIELD);
-        SNIPPETS_PARSER.declareInt(optionalConstructorArg(), MAX_SIZE_FIELD);
+        SNIPPETS_PARSER.declareInt(optionalConstructorArg(), NUM_SNIPPETS_FIELD);
 
         RetrieverBuilder.declareBaseParserFields(PARSER);
     }
@@ -235,11 +232,8 @@ protected void doToXContent(XContentBuilder builder, Params params) throws IOExc
         }
         if (snippets != null) {
             builder.startObject(SNIPPETS_FIELD.getPreferredName());
-            if (snippets.numFragments() != null) {
-                builder.field(NUM_FRAGMENTS_FIELD.getPreferredName(), snippets.numFragments());
-            }
-            if (snippets.maxSize() != null) {
-                builder.field(MAX_SIZE_FIELD.getPreferredName(), snippets.maxSize());
+            if (snippets.numSnippets() != null) {
+                builder.field(NUM_SNIPPETS_FIELD.getPreferredName(), snippets.numSnippets());
             }
             builder.endObject();
         }