[WIP] Support extract_snippets function in ES|QL

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

@@ -123,7 +123,7 @@ public SearchHit(int nestedTopDocId, String id, NestedIdentity nestedIdentity) {
         this(nestedTopDocId, id, nestedIdentity, null);
     }
 
-    private SearchHit(int nestedTopDocId, String id, NestedIdentity nestedIdentity, @Nullable RefCounted refCounted) {
+    public SearchHit(int nestedTopDocId, String id, NestedIdentity nestedIdentity, @Nullable RefCounted refCounted) {
         this(
             nestedTopDocId,
             DEFAULT_SCORE,

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

@@ -279,6 +279,8 @@
  * Sets up things that can be done at search time like queries, aggregations, and suggesters.
  */
 public class SearchModule {
+    private static volatile Map<String, Highlighter> staticHighlighters = Map.of();
+
     public static final Setting<Integer> INDICES_MAX_CLAUSE_COUNT_SETTING = Setting.intSetting(
         "indices.query.bool.max_clause_count",
         4096,
@@ -921,6 +923,10 @@ private static Map<String, Highlighter> setupHighlighters(Settings settings, Lis
         return unmodifiableMap(highlighters.getRegistry());
     }
 
+    public static Map<String, Highlighter> getStaticHighlighters() {
+        return staticHighlighters;
+    }
+
     private void registerScoreFunctions(List<SearchPlugin> plugins) {
         // ScriptScoreFunctionBuilder has it own named writable because of a new script_score query
         namedWriteables.add(
@@ -1060,6 +1066,9 @@ private void registerFetchSubPhases(List<SearchPlugin> plugins) {
         registerFetchSubPhase(new HighlightPhase(highlighters));
         registerFetchSubPhase(new FetchScorePhase());
 
+        // Store highlighters in a static map for other plugins to access
+        staticHighlighters = Map.copyOf(highlighters);
+
         FetchPhaseConstructionContext context = new FetchPhaseConstructionContext(highlighters);
         registerFromPlugin(plugins, p -> p.getFetchSubPhases(context), this::registerFetchSubPhase);
     }

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

@@ -114,7 +114,7 @@ public HighlightField highlight(FieldHighlightContext fieldContext) throws IOExc
 
     CustomUnifiedHighlighter buildHighlighter(FieldHighlightContext fieldContext) {
         IndexSettings indexSettings = fieldContext.context.getSearchExecutionContext().getIndexSettings();
-        Encoder encoder = fieldContext.field.fieldOptions().encoder().equals("html")
+        Encoder encoder = "html".equals(fieldContext.field.fieldOptions().encoder())
             ? HighlightUtils.Encoders.HTML
             : HighlightUtils.Encoders.DEFAULT;
 

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

@@ -185,16 +185,16 @@ public Map<String, Object> options() {
             return options;
         }
 
-        static class Builder {
+        public static class Builder {
 
             private final FieldOptions fieldOptions = new FieldOptions();
 
-            Builder fragmentCharSize(int fragmentCharSize) {
+            public Builder fragmentCharSize(int fragmentCharSize) {
                 fieldOptions.fragmentCharSize = fragmentCharSize;
                 return this;
             }
 
-            Builder numberOfFragments(int numberOfFragments) {
+            public Builder numberOfFragments(int numberOfFragments) {
                 fieldOptions.numberOfFragments = numberOfFragments;
                 return this;
             }
@@ -209,17 +209,17 @@ Builder encoder(String encoder) {
                 return this;
             }
 
-            Builder preTags(String[] preTags) {
+            public Builder preTags(String[] preTags) {
                 fieldOptions.preTags = preTags;
                 return this;
             }
 
-            Builder postTags(String[] postTags) {
+            public Builder postTags(String[] postTags) {
                 fieldOptions.postTags = postTags;
                 return this;
             }
 
-            Builder scoreOrdered(boolean scoreOrdered) {
+            public Builder scoreOrdered(boolean scoreOrdered) {
                 fieldOptions.scoreOrdered = scoreOrdered;
                 return this;
             }
@@ -229,7 +229,7 @@ Builder highlightFilter(boolean highlightFilter) {
                 return this;
             }
 
-            Builder requireFieldMatch(boolean requireFieldMatch) {
+            public Builder requireFieldMatch(boolean requireFieldMatch) {
                 fieldOptions.requireFieldMatch = requireFieldMatch;
                 return this;
             }
@@ -269,7 +269,7 @@ Builder boundaryScannerLocale(Locale boundaryScannerLocale) {
                 return this;
             }
 
-            Builder highlightQuery(Query highlightQuery) {
+            public Builder highlightQuery(Query highlightQuery) {
                 fieldOptions.highlightQuery = highlightQuery;
                 return this;
             }
@@ -294,7 +294,7 @@ Builder options(Map<String, Object> options) {
                 return this;
             }
 
-            FieldOptions build() {
+            public FieldOptions build() {
                 return fieldOptions;
             }
 

x-pack/plugin/esql/compute/src/main/java/module-info.java

@@ -21,6 +21,7 @@
     requires org.elasticsearch.geo;
     requires org.elasticsearch.xcore;
     requires hppc;
+    requires org.apache.lucene.highlighter;
 
     exports org.elasticsearch.compute;
     exports org.elasticsearch.compute.aggregation;

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

@@ -0,0 +1,184 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.compute.lucene;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.Scorable;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.DriverContext;
+import org.elasticsearch.compute.operator.EvalOperator;
+import org.elasticsearch.index.fieldvisitor.LeafStoredFieldLoader;
+import org.elasticsearch.index.fieldvisitor.StoredFieldLoader;
+import org.elasticsearch.index.mapper.MappedFieldType;
+import org.elasticsearch.index.mapper.SourceLoader;
+import org.elasticsearch.search.SearchHit;
+import org.elasticsearch.search.fetch.FetchContext;
+import org.elasticsearch.search.fetch.FetchSubPhase;
+import org.elasticsearch.search.fetch.subphase.highlight.DefaultHighlighter;
+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.Highlighter;
+import org.elasticsearch.search.fetch.subphase.highlight.SearchHighlightContext;
+import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.search.lookup.Source;
+import org.elasticsearch.xcontent.Text;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Supplier;
+
+import static org.elasticsearch.core.RefCounted.ALWAYS_REFERENCED;
+
+public class HighlighterExpressionEvaluator extends LuceneQueryEvaluator<BytesRefBlock.Builder>
+    implements
+        EvalOperator.ExpressionEvaluator {
+
+    private final String fieldName;
+    private final Integer numFragments;
+    private final Integer fragmentLength;
+    private final SearchContext searchContext;
+    private final Map<String, Highlighter> highlighters;
+
+    HighlighterExpressionEvaluator(
+        BlockFactory blockFactory,
+        ShardConfig[] shardConfigs,
+        String fieldName,
+        Integer numFragments,
+        Integer fragmentLength,
+        SearchContext searchContext,
+        Map<String, Highlighter> highlighters
+    ) {
+        super(blockFactory, shardConfigs);
+        this.fieldName = fieldName;
+        this.numFragments = numFragments;
+        this.fragmentLength = fragmentLength;
+        this.searchContext = searchContext;
+        this.highlighters = highlighters;
+    }
+
+    @Override
+    protected ScoreMode scoreMode() {
+        return ScoreMode.COMPLETE;
+    }
+
+    @Override
+    protected Block createNoMatchBlock(BlockFactory blockFactory, int size) {
+        return blockFactory.newConstantNullBlock(size);
+    }
+
+    @Override
+    protected BytesRefBlock.Builder createBlockBuilder(BlockFactory blockFactory, int size) {
+        return blockFactory.newBytesRefBlockBuilder(size * numFragments);
+    }
+
+    @Override
+    protected void appendMatch(BytesRefBlock.Builder builder, Scorable scorer, int docId, LeafReaderContext leafReaderContext, Query query)
+        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, null, null, ALWAYS_REFERENCED);
+        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());
+
+        FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext(searchHit, leafReaderContext, docId, Map.of(), source, null);
+        FieldHighlightContext highlightContext = new FieldHighlightContext(
+            fieldName,
+            field,
+            fieldType,
+            fetchContext,
+            hitContext,
+            query,
+            new HashMap<>()
+        );
+        HighlightField highlight = highlighter.highlight(highlightContext);
+
+        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();
+            }
+        }
+    }
+
+    private static Supplier<Source> lazyStoredSourceLoader(LeafReaderContext ctx, int doc) {
+        return () -> {
+            StoredFieldLoader rootLoader = StoredFieldLoader.create(true, Collections.emptySet());
+            try {
+                LeafStoredFieldLoader leafRootLoader = rootLoader.getLoader(ctx, null);
+                leafRootLoader.advanceTo(doc);
+                return Source.fromBytes(leafRootLoader.source());
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+        };
+    }
+
+    @Override
+    protected void appendNoMatch(BytesRefBlock.Builder builder) {
+        builder.appendNull();
+    }
+
+    @Override
+    public Block eval(Page page) {
+        return executeQuery(page);
+    }
+
+    public record Factory(
+        ShardConfig[] shardConfigs,
+        String fieldName,
+        Integer numFragments,
+        Integer fragmentSize,
+        SearchContext searchContext,
+        Map<String, Highlighter> highlighters
+    ) implements EvalOperator.ExpressionEvaluator.Factory {
+        @Override
+        public EvalOperator.ExpressionEvaluator get(DriverContext context) {
+            return new HighlighterExpressionEvaluator(
+                context.blockFactory(),
+                shardConfigs,
+                fieldName,
+                numFragments,
+                fragmentSize,
+                searchContext,
+                highlighters
+            );
+        }
+    }
+}