Wrapper for MatchAllDocsQuery

Author: piergm

server/src/internalClusterTest/java/org/elasticsearch/search/fetch/subphase/highlight/HighlighterSearchIT.java

@@ -3640,7 +3640,7 @@ public void testConstantKeywordAndOtherFieldsMatchMultipleHitsHighlighted() thro
         XContentBuilder mappings = prepareConstantKeywordMappings(constantKeywordFieldName, constantValue);
         assertAcked(prepareCreate(index).setMapping(mappings));
 
-        XContentBuilder firstDocument = jsonBuilder().startObject().field("foo", "bar").endObject();
+        XContentBuilder firstDocument = jsonBuilder().startObject().field("foo", constantValue).endObject();
         XContentBuilder secondDocument = jsonBuilder().startObject().field("foo", constantValue).endObject();
         saveDocumentIntoIndex(index, "1", firstDocument);
         saveDocumentIntoIndex(index, "2", secondDocument);
@@ -3649,8 +3649,9 @@ public void testConstantKeywordAndOtherFieldsMatchMultipleHitsHighlighted() thro
 
         assertNoFailures(search);
         assertHighlight(search, 0, constantKeywordFieldName, 0, 1, equalTo("<em>%s</em>".formatted(constantValue)));
-        assertHighlight(search, 1, "foo", 0, 1, equalTo("<em>%s</em>".formatted(constantValue)));
+        assertHighlight(search, 0, "foo", 0, 1, equalTo("<em>%s</em>".formatted(constantValue)));
         assertHighlight(search, 1, constantKeywordFieldName, 0, 1, equalTo("<em>%s</em>".formatted(constantValue)));
+        assertHighlight(search, 1, "foo", 0, 1, equalTo("<em>%s</em>".formatted(constantValue)));
     }
 
     public void testDoubleConstantKeywordJustOneHighlighted() throws IOException {

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

@@ -23,6 +23,7 @@
 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;
@@ -35,6 +36,10 @@ 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

@@ -54,7 +54,7 @@ private static String valueToString(Object value) {
     public final Query termQuery(Object value, SearchExecutionContext context) {
         String pattern = valueToString(value);
         if (matches(pattern, false, context)) {
-            return Queries.newMatchAllQuery();
+            return Queries.newMatchAllQueryWrapper(pattern);
         } else {
             return new MatchNoDocsQuery();
         }
@@ -64,7 +64,7 @@ public final Query termQuery(Object value, SearchExecutionContext context) {
     public final Query termQueryCaseInsensitive(Object value, SearchExecutionContext context) {
         String pattern = valueToString(value);
         if (matches(pattern, true, context)) {
-            return Queries.newMatchAllQuery();
+            return Queries.newMatchAllQueryWrapper(pattern);
         } else {
             return new MatchNoDocsQuery();
         }
@@ -76,7 +76,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.newMatchAllQuery();
+                return Queries.newMatchAllQueryWrapper(pattern);
             }
         }
         return new MatchNoDocsQuery();
@@ -91,7 +91,7 @@ public final Query prefixQuery(
     ) {
         String pattern = prefix + "*";
         if (matches(pattern, caseInsensitive, context)) {
-            return Queries.newMatchAllQuery();
+            return Queries.newMatchAllQueryWrapper(pattern);
         } else {
             return new MatchNoDocsQuery();
         }
@@ -105,7 +105,7 @@ public final Query wildcardQuery(
         SearchExecutionContext context
     ) {
         if (matches(value, caseInsensitive, context)) {
-            return Queries.newMatchAllQuery();
+            return Queries.newMatchAllQueryWrapper(value);
         } else {
             return new MatchNoDocsQuery();
         }

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

@@ -0,0 +1,55 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.index.query;
+
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.Weight;
+
+public final class MatchAllDocumentQueryWrapper extends Query {
+    private final String pattern;
+    private final MatchAllDocsQuery matchAllDocsQuery;
+
+    public MatchAllDocumentQueryWrapper(String pattern) {
+        this.matchAllDocsQuery = new MatchAllDocsQuery();
+        this.pattern = pattern;
+    }
+
+    public String getPattern() {
+        return pattern;
+    }
+
+    @Override
+    public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) {
+        return matchAllDocsQuery.createWeight(searcher, scoreMode, boost);
+    }
+
+    @Override
+    public String toString(String field) {
+        return matchAllDocsQuery.toString();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        return matchAllDocsQuery.equals(o);
+    }
+
+    @Override
+    public int hashCode() {
+        return matchAllDocsQuery.hashCode();
+    }
+
+    @Override
+    public void visit(QueryVisitor visitor) {
+        visitor.visitLeaf(matchAllDocsQuery);
+    }
+}

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

@@ -11,14 +11,15 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.DisjunctionMaxQuery;
-import org.apache.lucene.search.MatchAllDocsQuery;
 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.search.fetch.FetchContext;
 import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.fetch.FetchSubPhaseProcessor;
@@ -174,12 +175,14 @@ private Query getHighlighterQuery(FetchContext context, SearchHighlightContext.F
     }
 
     private Query getHighlighterQueryForConstantFieldType(FetchContext context, MappedFieldType fieldType) {
-        if (context.query() instanceof MatchAllDocsQuery) {
-            return new TermQuery(new Term(fieldType.name(), getValueForConstantFieldType(context, fieldType)));
-        } else if (context.query() instanceof DisjunctionMaxQuery disjunctionMaxQuery) {
-            for (Query anyDisjunctQuery : disjunctionMaxQuery.getDisjuncts()) {
-                if (anyDisjunctQuery instanceof MatchAllDocsQuery) {
-                    return new TermQuery(new Term(fieldType.name(), getValueForConstantFieldType(context, fieldType)));
+        if (context.query() instanceof MatchAllDocumentQueryWrapper persist) {
+            return new TermQuery(new Term(fieldType.name(), persist.getPattern()));
+        } else if (context.query() instanceof DisjunctionMaxQuery changeName) {
+            for (Query query : changeName.getDisjuncts()) {
+                if (query instanceof MatchAllDocumentQueryWrapper anyPersist) {
+                    if (Regex.simpleMatch(anyPersist.getPattern(), getValueForConstantFieldType(context, fieldType))) {
+                        return new TermQuery(new Term(fieldType.name(), getValueForConstantFieldType(context, fieldType)));
+                    }
                 }
             }
         }