Rewrite wrapper queries to match_none if possible. (#55271)

Queries like script_score wrap a query and modify its score. If the inner query
rewrites to match_none, then the entire query can rewrite to match_none. This
lets us detect that certain shards can be skipped during the 'can match' phase.

This was a simple change that seemed like it would help in some cases. But it
will likely not have a huge impact, since in many use cases where the 'can
match' phase is helpful, the search is not sorted by score.

Author: jtibshirani

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

@@ -219,6 +219,10 @@ protected boolean doEquals(BoostingQueryBuilder other) {
     @Override
     protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws IOException {
         QueryBuilder positiveQuery = this.positiveQuery.rewrite(queryRewriteContext);
+        if (positiveQuery instanceof MatchNoneQueryBuilder) {
+            return positiveQuery;
+        }
+
         QueryBuilder negativeQuery = this.negativeQuery.rewrite(queryRewriteContext);
         if (positiveQuery != this.positiveQuery || negativeQuery != this.negativeQuery) {
             BoostingQueryBuilder newQueryBuilder = new BoostingQueryBuilder(positiveQuery, negativeQuery);

server/src/main/java/org/elasticsearch/index/query/functionscore/FunctionScoreQueryBuilder.java

@@ -36,6 +36,7 @@
 import org.elasticsearch.index.query.AbstractQueryBuilder;
 import org.elasticsearch.index.query.InnerHitContextBuilder;
 import org.elasticsearch.index.query.MatchAllQueryBuilder;
+import org.elasticsearch.index.query.MatchNoneQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.QueryShardContext;
@@ -405,6 +406,10 @@ public FilterFunctionBuilder rewrite(QueryRewriteContext context) throws IOExcep
     @Override
     protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws IOException {
         QueryBuilder queryBuilder = this.query.rewrite(queryRewriteContext);
+        if (queryBuilder instanceof MatchNoneQueryBuilder) {
+            return queryBuilder;
+        }
+
         FilterFunctionBuilder[] rewrittenBuilders = new FilterFunctionBuilder[this.filterFunctionBuilders.length];
         boolean rewritten = false;
         for (int i = 0; i < rewrittenBuilders.length; i++) {

server/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreQueryBuilder.java

@@ -31,6 +31,7 @@
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
 import org.elasticsearch.index.query.InnerHitContextBuilder;
+import org.elasticsearch.index.query.MatchNoneQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.QueryShardContext;
@@ -187,6 +188,10 @@ protected Query doToQuery(QueryShardContext context) throws IOException {
     @Override
     protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws IOException {
         QueryBuilder newQuery = this.query.rewrite(queryRewriteContext);
+        if (newQuery instanceof MatchNoneQueryBuilder) {
+            return newQuery;
+        }
+
         if (newQuery != query) {
             ScriptScoreQueryBuilder newQueryBuilder = new ScriptScoreQueryBuilder(newQuery, script);
             if (minScore != null) {

server/src/test/java/org/elasticsearch/index/query/BoostingQueryBuilderTests.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.index.query;
 
 import org.apache.lucene.queries.function.FunctionScoreQuery;
+import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.elasticsearch.test.AbstractQueryTestCase;
 
@@ -44,6 +45,8 @@ protected void doAssertLuceneQuery(BoostingQueryBuilder queryBuilder, Query quer
         Query negative = queryBuilder.negativeQuery().rewrite(context).toQuery(context);
         if (positive == null || negative == null) {
             assertThat(query, nullValue());
+        } else if (positive instanceof MatchNoDocsQuery) {
+            assertThat(query, instanceOf(MatchNoDocsQuery.class));
         } else {
             assertThat(query, instanceOf(FunctionScoreQuery.class));
         }
@@ -91,9 +94,9 @@ public void testFromJson() throws IOException {
 
     public void testRewrite() throws IOException {
         QueryBuilder positive = randomBoolean() ? new MatchAllQueryBuilder() :
-            new WrapperQueryBuilder(new TermQueryBuilder("pos", "bar").toString());
+            new WrapperQueryBuilder(new TermQueryBuilder(KEYWORD_FIELD_NAME, "bar").toString());
         QueryBuilder negative = randomBoolean() ? new MatchAllQueryBuilder() :
-            new WrapperQueryBuilder(new TermQueryBuilder("neg", "bar").toString());
+            new WrapperQueryBuilder(new TermQueryBuilder(TEXT_FIELD_NAME, "bar").toString());
         BoostingQueryBuilder qb = new BoostingQueryBuilder(positive, negative);
         QueryBuilder rewrite = qb.rewrite(createShardContext());
         if (positive instanceof MatchAllQueryBuilder && negative instanceof MatchAllQueryBuilder) {
@@ -104,6 +107,14 @@ public void testRewrite() throws IOException {
         }
     }
 
+    public void testRewriteToMatchNone() throws IOException {
+        BoostingQueryBuilder builder = new BoostingQueryBuilder(
+            new TermQueryBuilder("unmapped_field", "value"),
+            new TermQueryBuilder(KEYWORD_FIELD_NAME, "other_value"));
+        QueryBuilder rewrite = builder.rewrite(createShardContext());
+        assertThat(rewrite, instanceOf(MatchNoneQueryBuilder.class));
+    }
+
     @Override
     public void testMustRewrite() throws IOException {
         QueryShardContext context = createShardContext();

server/src/test/java/org/elasticsearch/index/query/ScriptScoreQueryBuilderTests.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.index.query;
 
+import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.lucene.search.function.ScriptScoreQuery;
@@ -51,7 +52,12 @@ protected ScriptScoreQueryBuilder doCreateTestQueryBuilder() {
 
     @Override
     protected void doAssertLuceneQuery(ScriptScoreQueryBuilder queryBuilder, Query query, QueryShardContext context) throws IOException {
-        assertThat(query, instanceOf(ScriptScoreQuery.class));
+        Query wrappedQuery = queryBuilder.query().rewrite(context).toQuery(context);
+        if (wrappedQuery instanceof MatchNoDocsQuery) {
+            assertThat(query, instanceOf(MatchNoDocsQuery.class));
+        } else {
+            assertThat(query, instanceOf(ScriptScoreQuery.class));
+        }
     }
 
     public void testFromJson() throws IOException {
@@ -92,7 +98,10 @@ public void testIllegalArguments() {
      */
     @Override
     public void testCacheability() throws IOException {
-        ScriptScoreQueryBuilder queryBuilder = createTestQueryBuilder();
+        Script script = new Script(ScriptType.INLINE, MockScriptEngine.NAME, "1", Collections.emptyMap());
+        ScriptScoreQueryBuilder queryBuilder = new ScriptScoreQueryBuilder(
+            new TermQueryBuilder(KEYWORD_FIELD_NAME, "value"), script);
+
         QueryShardContext context = createShardContext();
         QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
         assertNotNull(rewriteQuery.toQuery(context));
@@ -112,6 +121,13 @@ public void testMustRewrite() throws IOException {
         assertEquals("Rewrite first", e.getMessage());
     }
 
+    public void testRewriteToMatchNone() throws IOException {
+        Script script = new Script(ScriptType.INLINE, MockScriptEngine.NAME, "1", Collections.emptyMap());
+        ScriptScoreQueryBuilder builder = new ScriptScoreQueryBuilder(new TermQueryBuilder("unmapped_field", "value"), script);
+        QueryBuilder rewrite = builder.rewrite(createShardContext());
+        assertThat(rewrite, instanceOf(MatchNoneQueryBuilder.class));
+    }
+
     public void testDisallowExpensiveQueries() {
         QueryShardContext queryShardContext = mock(QueryShardContext.class);
         when(queryShardContext.allowExpensiveQueries()).thenReturn(false);

server/src/test/java/org/elasticsearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java

@@ -20,9 +20,9 @@
 package org.elasticsearch.index.query.functionscore;
 
 import com.fasterxml.jackson.core.JsonParseException;
-
 import org.apache.lucene.index.Term;
 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.ParsingException;
@@ -40,6 +40,7 @@
 import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.index.mapper.SeqNoFieldMapper;
 import org.elasticsearch.index.query.MatchAllQueryBuilder;
+import org.elasticsearch.index.query.MatchNoneQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.query.RandomQueryBuilder;
@@ -54,6 +55,7 @@
 import org.elasticsearch.search.MultiValueMode;
 import org.elasticsearch.test.AbstractQueryTestCase;
 import org.elasticsearch.test.TestGeoShapeFieldMapperPlugin;
+import org.hamcrest.CoreMatchers;
 import org.hamcrest.Matcher;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
@@ -77,7 +79,6 @@
 import static org.elasticsearch.index.query.functionscore.ScoreFunctionBuilders.weightFactorFunction;
 import static org.hamcrest.Matchers.closeTo;
 import static org.hamcrest.Matchers.containsString;
-import static org.hamcrest.Matchers.either;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.nullValue;
@@ -254,7 +255,12 @@ private static DecayFunctionBuilder<?> createRandomDecayFunction() {
 
     @Override
     protected void doAssertLuceneQuery(FunctionScoreQueryBuilder queryBuilder, Query query, QueryShardContext context) throws IOException {
-        assertThat(query, either(instanceOf(FunctionScoreQuery.class)).or(instanceOf(FunctionScoreQuery.class)));
+        Query wrappedQuery = queryBuilder.query().rewrite(context).toQuery(context);
+        if (wrappedQuery instanceof MatchNoDocsQuery) {
+            assertThat(query, CoreMatchers.instanceOf(MatchNoDocsQuery.class));
+        } else {
+            assertThat(query, CoreMatchers.instanceOf(FunctionScoreQuery.class));
+        }
     }
 
     public void testIllegalArguments() {
@@ -674,11 +680,23 @@ public void testRewriteWithFunction() throws IOException {
         assertSame(rewrite.filterFunctionBuilders()[1].getFilter(), secondFunction);
     }
 
+    public void testRewriteToMatchNone() throws IOException {
+        FunctionScoreQueryBuilder functionScoreQueryBuilder =
+            new FunctionScoreQueryBuilder(new TermQueryBuilder("unmapped_field", "value"))
+                .boostMode(CombineFunction.REPLACE)
+                .scoreMode(FunctionScoreQuery.ScoreMode.SUM)
+                .setMinScore(1)
+                .maxBoost(100);
+
+        QueryBuilder rewrite = functionScoreQueryBuilder.rewrite(createShardContext());
+        assertThat(rewrite, instanceOf(MatchNoneQueryBuilder.class));
+    }
+
     /**
      * Please see https://github.com/elastic/elasticsearch/issues/35123 for context.
      */
     public void testSingleScriptFunction() throws IOException {
-        QueryBuilder queryBuilder = RandomQueryBuilder.createQuery(random());
+        QueryBuilder queryBuilder = termQuery(KEYWORD_FIELD_NAME, "value");
         ScoreFunctionBuilder<ScriptScoreFunctionBuilder> functionBuilder = new ScriptScoreFunctionBuilder(
             new Script(ScriptType.INLINE, MockScriptEngine.NAME, "1", Collections.emptyMap()));