cleaner implementation

Author: piergm

server/src/main/java/org/elasticsearch/common/lucene/search/Queries.java

@@ -23,7 +23,6 @@
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.mapper.NestedPathFieldMapper;
 import org.elasticsearch.index.mapper.SeqNoFieldMapper;
-import org.elasticsearch.index.query.MatchAllDocumentQueryWrapper;
 
 import java.util.Collection;
 import java.util.Collections;
@@ -36,10 +35,6 @@ public static Query newMatchAllQuery() {
         return new MatchAllDocsQuery();
     }
 
-    public static Query newMatchAllQueryWrapper(String pattern) {
-        return new MatchAllDocumentQueryWrapper(pattern);
-    }
-
     /** Return a query that matches no document. */
     public static Query newMatchNoDocsQuery(String reason) {
         return new MatchNoDocsQuery(reason);

server/src/main/java/org/elasticsearch/index/mapper/ConstantFieldType.java

@@ -62,7 +62,7 @@ public final Query termQuery(Object value, SearchExecutionContext context) {
     public final Query internalTermQuery(Object value, QueryRewriteContext context) {
         String pattern = valueToString(value);
         if (matches(pattern, false, context)) {
-            return Queries.newMatchAllQueryWrapper(pattern);
+            return Queries.newMatchAllQuery();
         } else {
             return new MatchNoDocsQuery();
         }
@@ -76,7 +76,7 @@ public final Query termQueryCaseInsensitive(Object value, SearchExecutionContext
     public final Query internalTermQueryCaseInsensitive(Object value, QueryRewriteContext context) {
         String pattern = valueToString(value);
         if (matches(pattern, true, context)) {
-            return Queries.newMatchAllQueryWrapper(pattern);
+            return Queries.newMatchAllQuery();
         } else {
             return new MatchNoDocsQuery();
         }
@@ -88,7 +88,7 @@ public final Query termsQuery(Collection<?> values, SearchExecutionContext conte
             String pattern = valueToString(value);
             if (matches(pattern, false, context)) {
                 // `terms` queries are a disjunction, so one matching term is enough
-                return Queries.newMatchAllQueryWrapper(pattern);
+                return Queries.newMatchAllQuery();
             }
         }
         return new MatchNoDocsQuery();
@@ -103,7 +103,7 @@ public final Query prefixQuery(
     ) {
         String pattern = prefix + "*";
         if (matches(pattern, caseInsensitive, context)) {
-            return Queries.newMatchAllQueryWrapper(pattern);
+            return Queries.newMatchAllQuery();
         } else {
             return new MatchNoDocsQuery();
         }
@@ -117,7 +117,7 @@ public final Query wildcardQuery(
         SearchExecutionContext context
     ) {
         if (matches(value, caseInsensitive, context)) {
-            return Queries.newMatchAllQueryWrapper(value);
+            return Queries.newMatchAllQuery();
         } else {
             return new MatchNoDocsQuery();
         }

server/src/main/java/org/elasticsearch/index/query/MatchAllDocumentQueryWrapper.java

@@ -26,6 +26,18 @@ public interface QueryBuilder extends VersionedNamedWriteable, ToXContentObject,
      */
     Query toQuery(SearchExecutionContext context) throws IOException;
 
+    /**
+     * Converts this QueryBuilder to a lucene {@link Query}.
+     * Returns {@code null} if this query should be ignored in the context of highlighting.
+     *
+     * @param fieldName name of the field we want to highlight
+     * @return the {@link Query} or {@code null} if this query should have standard behaviour
+     * for highlighting.
+     */
+    default Query toDoHighlight(String fieldName) {
+        return null;
+    }
+
     /**
      * Sets the arbitrary name to be assigned to the query (see named queries).
      * Implementers should return the concrete type of the

server/src/main/java/org/elasticsearch/index/query/QueryBuilder.java

@@ -8,9 +8,11 @@
 
 package org.elasticsearch.index.query;
 
+import org.apache.lucene.index.Term;
 import org.apache.lucene.search.BoostQuery;
 import org.apache.lucene.search.FuzzyQuery;
 import org.apache.lucene.search.Query;
+import org.apache.lucene.search.WildcardQuery;
 import org.apache.lucene.util.automaton.Operations;
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.common.ParsingException;
@@ -970,4 +972,9 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
     public TransportVersion getMinimalSupportedVersion() {
         return TransportVersion.ZERO;
     }
+
+    @Override
+    public Query toDoHighlight(String fieldName) {
+        return new WildcardQuery(new Term(fieldName, queryString));
+    }
 }

server/src/main/java/org/elasticsearch/index/query/QueryStringQueryBuilder.java

@@ -90,6 +90,7 @@
 import org.elasticsearch.search.fetch.subphase.FetchFieldsContext;
 import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
+import org.elasticsearch.search.fetch.subphase.highlight.SearchHighlightContext;
 import org.elasticsearch.search.internal.AliasFilter;
 import org.elasticsearch.search.internal.InternalScrollSearchRequest;
 import org.elasticsearch.search.internal.LegacyReaderContext;
@@ -1284,7 +1285,9 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
         if (source.highlighter() != null) {
             HighlightBuilder highlightBuilder = source.highlighter();
             try {
-                context.highlight(highlightBuilder.build(searchExecutionContext));
+                SearchHighlightContext searchHighlightContext = highlightBuilder.build(searchExecutionContext);
+                searchHighlightContext.setOriginalQuery(source.query());
+                context.highlight(searchHighlightContext);
             } catch (IOException e) {
                 throw new SearchException(shardTarget, "failed to create SearchContextHighlighter", e);
             }

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

@@ -9,17 +9,12 @@
 package org.elasticsearch.search.fetch.subphase.highlight;
 
 import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.index.Term;
-import org.apache.lucene.search.DisjunctionMaxQuery;
-import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
-import org.apache.lucene.search.TermQuery;
-import org.elasticsearch.common.regex.Regex;
 import org.elasticsearch.index.mapper.ConstantFieldType;
 import org.elasticsearch.index.mapper.KeywordFieldMapper;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.TextFieldMapper;
-import org.elasticsearch.index.query.MatchAllDocumentQueryWrapper;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.search.fetch.FetchContext;
 import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.fetch.FetchSubPhaseProcessor;
@@ -146,7 +141,7 @@ private FieldContext contextBuilders(
                     sourceRequired = true;
                 }
 
-                Query highlightQuery = getHighlightQuery(context, field, fieldType);
+                Query highlightQuery = getHighlightQuery(highlightContext, field, fieldType);
 
                 builders.put(
                     fieldName,
@@ -166,25 +161,15 @@ private FieldContext contextBuilders(
         return new FieldContext(storedFieldsSpec, builders);
     }
 
-    private Query getHighlightQuery(FetchContext context, SearchHighlightContext.Field field, MappedFieldType fieldType) {
-        if (fieldType instanceof ConstantFieldType constantFieldType) {
-            return getHighlightQueryForConstantFieldType(context, constantFieldType);
+    private Query getHighlightQuery(
+        SearchHighlightContext highlightContext,
+        SearchHighlightContext.Field field,
+        MappedFieldType fieldType
+    ) {
+        if (fieldType instanceof ConstantFieldType) {
+            QueryBuilder originalQuery = highlightContext.getOriginalQuery();
+            return originalQuery.toDoHighlight(fieldType.name());
         }
         return field.fieldOptions().highlightQuery();
     }
-
-    private Query getHighlightQueryForConstantFieldType(FetchContext context, ConstantFieldType fieldType) {
-        if (context.query() instanceof MatchAllDocumentQueryWrapper matchAllDocumentQueryWrapper) {
-            return new TermQuery(new Term(fieldType.name(), matchAllDocumentQueryWrapper.getPattern()));
-        } else if (context.query() instanceof DisjunctionMaxQuery unionOfQueries) {
-            for (Query query : unionOfQueries.getDisjuncts()) {
-                if (query instanceof MatchAllDocumentQueryWrapper matchAllDocumentQueryWrapper) {
-                    if (Regex.simpleMatch(matchAllDocumentQueryWrapper.getPattern(), fieldType.getConstantValue(context))) {
-                        return new TermQuery(new Term(fieldType.name(), fieldType.getConstantValue(context)));
-                    }
-                }
-            }
-        }
-        return new MatchNoDocsQuery();
-    }
 }

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

@@ -10,6 +10,7 @@
 
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.util.Maps;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder.BoundaryScannerType;
 
 import java.util.Arrays;
@@ -20,7 +21,11 @@
 import java.util.Set;
 
 public class SearchHighlightContext {
+    public void setOriginalQuery(QueryBuilder originalQuery) {
+        this.originalQuery = originalQuery;
+    }
 
+    private QueryBuilder originalQuery;
     private final Map<String, Field> fields;
 
     public SearchHighlightContext(Collection<Field> fields) {
@@ -35,6 +40,10 @@ public Collection<Field> fields() {
         return fields.values();
     }
 
+    public QueryBuilder getOriginalQuery() {
+        return originalQuery;
+    }
+
     public static class Field {
         private final String field;
         private final FieldOptions fieldOptions;