Refactor highlighting logic into util class

Author: kderusso

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

@@ -0,0 +1,60 @@
+/*
+ * 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.fetch.subphase.highlight;
+
+import org.apache.lucene.search.Query;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.SearchExecutionContext;
+
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * Utility class for building highlighting queries for the purpose of extracting snippets.
+ */
+public class HighlightSnippetUtils {
+
+    public static SearchHighlightContext buildSearchHighlightContextForSnippets(
+        SearchExecutionContext searchExecutionContext,
+        String field,
+        int numSnippets,
+        int snippetCharLength,
+        QueryBuilder queryBuilder
+    ) throws IOException {
+        SearchHighlightContext.Field highlightField = buildFieldHighlightContextForSnippets(
+            searchExecutionContext,
+            field,
+            numSnippets,
+            snippetCharLength,
+            queryBuilder.toQuery(searchExecutionContext)
+        );
+        return new SearchHighlightContext(List.of(highlightField));
+    }
+
+    public static SearchHighlightContext.Field buildFieldHighlightContextForSnippets(
+        SearchExecutionContext searchExecutionContext,
+        String fieldName,
+        int numSnippets,
+        int snippetCharLength,
+        Query query
+    ) {
+        SearchHighlightContext.FieldOptions.Builder optionsBuilder = new SearchHighlightContext.FieldOptions.Builder();
+        optionsBuilder.numberOfFragments(numSnippets);
+        optionsBuilder.fragmentCharSize(snippetCharLength);
+        optionsBuilder.noMatchSize(snippetCharLength);
+        optionsBuilder.preTags(new String[] { "" });
+        optionsBuilder.postTags(new String[] { "" });
+        optionsBuilder.requireFieldMatch(false);
+        optionsBuilder.scoreOrdered(true);
+        optionsBuilder.highlightQuery(query);
+        return new SearchHighlightContext.Field(fieldName, optionsBuilder.build());
+    }
+
+}

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/HighlighterExpressionEvaluator.java

@@ -30,6 +30,7 @@
 import org.elasticsearch.search.fetch.subphase.highlight.FieldHighlightContext;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
+import org.elasticsearch.search.fetch.subphase.highlight.HighlightSnippetUtils;
 import org.elasticsearch.search.fetch.subphase.highlight.Highlighter;
 import org.elasticsearch.search.fetch.subphase.highlight.SearchHighlightContext;
 import org.elasticsearch.search.internal.SearchContext;
@@ -106,17 +107,13 @@ protected void appendMatch(BytesRefBlock.Builder builder, Scorable scorer, int d
         Source source = Source.lazy(lazyStoredSourceLoader(leafReaderContext, docId));
         Highlighter highlighter = highlighters.getOrDefault(fieldType.getDefaultHighlighter(), new DefaultHighlighter());
 
-        // TODO: Consolidate these options with the ones built in the text similarity reranker
-        SearchHighlightContext.FieldOptions.Builder optionsBuilder = new SearchHighlightContext.FieldOptions.Builder();
-        optionsBuilder.numberOfFragments(numFragments != null ? numFragments : HighlightBuilder.DEFAULT_NUMBER_OF_FRAGMENTS);
-        optionsBuilder.fragmentCharSize(fragmentLength != null ? fragmentLength : HighlightBuilder.DEFAULT_FRAGMENT_CHAR_SIZE);
-        optionsBuilder.preTags(new String[] { "" });
-        optionsBuilder.postTags(new String[] { "" });
-        optionsBuilder.requireFieldMatch(false);
-        optionsBuilder.scoreOrdered(true);
-        optionsBuilder.highlightQuery(query);
-        SearchHighlightContext.Field field = new SearchHighlightContext.Field(fieldName, optionsBuilder.build());
-
+        SearchHighlightContext.Field field = HighlightSnippetUtils.buildFieldHighlightContextForSnippets(
+            fetchContext.getSearchExecutionContext(),
+            fieldName,
+            numFragments != null ? numFragments : HighlightBuilder.DEFAULT_NUMBER_OF_FRAGMENTS,
+            fragmentLength != null ? fragmentLength : HighlightBuilder.DEFAULT_FRAGMENT_CHAR_SIZE,
+            query
+        );
         FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext(searchHit, leafReaderContext, docId, Map.of(), source, null);
         FieldHighlightContext highlightContext = new FieldHighlightContext(
             fieldName,

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

@@ -12,8 +12,8 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.SearchHits;
-import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
+import org.elasticsearch.search.fetch.subphase.highlight.HighlightSnippetUtils;
 import org.elasticsearch.search.fetch.subphase.highlight.SearchHighlightContext;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.rank.RankShardResult;
@@ -73,20 +73,17 @@ public RankShardResult doBuildRankFeatureShardResult(SearchHits hits, int shardI
     public void prepareForFetch(SearchContext context) {
         if (snippetRankInput != null) {
             try {
-                HighlightBuilder highlightBuilder = new HighlightBuilder();
-                highlightBuilder.highlightQuery(snippetRankInput.snippetQueryBuilder());
-                // Stripping pre/post tags as they're not useful for snippet creation
-                highlightBuilder.field(field).preTags("").postTags("");
-                // Return highest scoring fragments
-                highlightBuilder.order(HighlightBuilder.Order.SCORE);
                 int numSnippets = snippetRankInput.numSnippets() != null ? snippetRankInput.numSnippets() : DEFAULT_NUM_SNIPPETS;
-                highlightBuilder.numOfFragments(numSnippets);
                 // Rely on the model to determine the fragment size
                 int tokenSizeLimit = snippetRankInput.tokenSizeLimit();
                 int fragmentSize = tokenSizeLimit * TOKEN_SIZE_LIMIT_MULTIPLIER;
-                highlightBuilder.fragmentSize(fragmentSize);
-                highlightBuilder.noMatchSize(fragmentSize);
-                SearchHighlightContext searchHighlightContext = highlightBuilder.build(context.getSearchExecutionContext());
+                SearchHighlightContext searchHighlightContext = HighlightSnippetUtils.buildSearchHighlightContextForSnippets(
+                    context.getSearchExecutionContext(),
+                    field,
+                    numSnippets,
+                    fragmentSize,
+                    snippetRankInput.snippetQueryBuilder()
+                );
                 context.highlight(searchHighlightContext);
             } catch (IOException e) {
                 throw new RuntimeException("Failed to generate snippet request", e);