Workaround for rewrite

Author: kderusso

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

@@ -60,10 +60,6 @@ public boolean canHighlight(MappedFieldType fieldType) {
     @Override
     public HighlightField highlight(FieldHighlightContext fieldContext) throws IOException {
         @SuppressWarnings("unchecked")
-        // Map<String, CustomUnifiedHighlighter> cache = (Map<String, CustomUnifiedHighlighter>) fieldContext.cache.computeIfAbsent(
-        // UnifiedHighlighter.class.getName(),
-        // k -> new HashMap<>()
-        // );
         Map<String, CustomUnifiedHighlighter> cache = (Map<String, CustomUnifiedHighlighter>) fieldContext.cache.getOrDefault(
             UnifiedHighlighter.class.getName(),
             new HashMap<>()

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

@@ -86,6 +86,22 @@ protected void appendMatch(BytesRefBlock.Builder builder, Scorable scorer, int d
         throws IOException {
 
         // TODO: Can we build a custom highlighter directly here, so we don't have to rely on fetch phase classes?
+
+        // Create a source loader for highlighter use
+        SourceLoader sourceLoader = searchContext.newSourceLoader(null);
+        FetchContext fetchContext = new FetchContext(searchContext, sourceLoader);
+        MappedFieldType fieldType = searchContext.getSearchExecutionContext().getFieldType(fieldName);
+        SearchHit searchHit = new SearchHit(docId);
+        Source source = Source.lazy(lazyStoredSourceLoader(leafReaderContext, docId));
+        String defaultHighlighter = fieldType.getDefaultHighlighter();
+
+        Highlighter highlighter;
+        // if (SemanticTextHighlighter.NAME.equals(defaultHighlighter)) {
+        // highlighter = new SemanticTextHighlighter();
+        // } else {
+        highlighter = new DefaultHighlighter();
+        // }
+
         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);
@@ -94,12 +110,6 @@ protected void appendMatch(BytesRefBlock.Builder builder, Scorable scorer, int d
         optionsBuilder.requireFieldMatch(false);
         optionsBuilder.scoreOrdered(true);
         SearchHighlightContext.Field field = new SearchHighlightContext.Field(fieldName, optionsBuilder.build());
-        // Create a source loader for highlighter use
-        SourceLoader sourceLoader = searchContext.newSourceLoader(null);
-        FetchContext fetchContext = new FetchContext(searchContext, sourceLoader);
-        MappedFieldType fieldType = searchContext.getSearchExecutionContext().getFieldType(fieldName);
-        SearchHit searchHit = new SearchHit(docId);
-        Source source = Source.lazy(lazyStoredSourceLoader(leafReaderContext, docId));
 
         FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext(searchHit, leafReaderContext, docId, Map.of(), source, null);
         FieldHighlightContext highlightContext = new FieldHighlightContext(
@@ -111,18 +121,19 @@ protected void appendMatch(BytesRefBlock.Builder builder, Scorable scorer, int d
             query,
             new HashMap<>()
         );
-        Highlighter highlighter = new DefaultHighlighter();
         HighlightField highlight = highlighter.highlight(highlightContext);
 
-        boolean multivalued = highlight.fragments().length > 1;
-        if (multivalued) {
-            builder.beginPositionEntry();
-        }
-        for (Text highlightText : highlight.fragments()) {
-            builder.appendBytesRef(new BytesRef(highlightText.bytes().bytes()));
-        }
-        if (multivalued) {
-            builder.endPositionEntry();
+        if (highlight != null) {
+            boolean multivalued = highlight.fragments().length > 1;
+            if (multivalued) {
+                builder.beginPositionEntry();
+            }
+            for (Text highlightText : highlight.fragments()) {
+                builder.appendBytesRef(new BytesRef(highlightText.bytes().bytes()));
+            }
+            if (multivalued) {
+                builder.endPositionEntry();
+            }
         }
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/QueryBuilderResolver.java

@@ -90,7 +90,8 @@ public FunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOException {
                         ? translationAware.asQuery(LucenePushdownPredicates.DEFAULT, TranslatorHandler.TRANSLATOR_HANDLER).toQueryBuilder()
                         : builder;
                     try {
-                        builder = builder.rewrite(ctx);
+                        // builder = builder.rewrite(ctx);
+                        builder = Rewriteable.rewrite(builder, ctx);
                     } catch (IOException e) {
                         exceptionHolder.setIfAbsent(e);
                     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/ExtractSnippets.java

@@ -15,6 +15,7 @@
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
 import org.elasticsearch.index.query.MatchQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.Rewriteable;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
 import org.elasticsearch.search.fetch.subphase.highlight.SearchHighlightContext;
@@ -217,7 +218,9 @@ public ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) {
                 // TODO: Reduce duplication between this method and TextSimilarityRerankingRankFeaturePhaseRankShardContext#prepareForFetch
                 HighlightBuilder highlightBuilder = new HighlightBuilder();
                 if (queryBuilder != null) {
-                    highlightBuilder.highlightQuery(queryBuilder);
+                    // TODO validate this works and determine why this is not working in query builder resolver
+                    QueryBuilder rewritten = Rewriteable.rewrite(queryBuilder, searchExecutionContext);
+                    highlightBuilder.highlightQuery(rewritten);
                 }
                 // Stripping pre/post tags as they're not useful for snippet creation
                 highlightBuilder.field(field.sourceText()).preTags("").postTags("");