inner_hits: Ensure that that InnerHitBuilder uses rewritten queries

If a nested, has_child or has_parent query's inner query gets rewritten then the InnerHitBuilder should use that rewritten form too, otherwise this can cause exceptions in a later phase.

Also fixes a bug that HasChildQueryBuilder's rewrite method overwrites max_children with min_children value.

Closes #19353

core/src/main/java/org/elasticsearch/index/query/HasChildQueryBuilder.java

@@ -472,9 +472,10 @@ public class HasChildQueryBuilder extends AbstractQueryBuilder<HasChildQueryBuil
 
     @Override
     protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws IOException {
-        QueryBuilder rewrite = query.rewrite(queryRewriteContext);
+        QueryBuilder rewrittenQuery = query.rewrite(queryRewriteContext);
-        if (rewrite != query) {
+        if (rewrittenQuery != query) {
-            return new HasChildQueryBuilder(type, rewrite, minChildren, minChildren, scoreMode, innerHitBuilder);
+            InnerHitBuilder rewrittenInnerHit = InnerHitBuilder.rewrite(innerHitBuilder, rewrittenQuery);
+            return new HasChildQueryBuilder(type, rewrittenQuery, minChildren, maxChildren, scoreMode, rewrittenInnerHit);
         }
         return this;
     }

core/src/main/java/org/elasticsearch/index/query/HasParentQueryBuilder.java

@@ -309,9 +309,10 @@ public class HasParentQueryBuilder extends AbstractQueryBuilder<HasParentQueryBu
 
     @Override
     protected QueryBuilder doRewrite(QueryRewriteContext queryShardContext) throws IOException {
-        QueryBuilder rewrite = query.rewrite(queryShardContext);
+        QueryBuilder rewrittenQuery = query.rewrite(queryShardContext);
-        if (rewrite != query) {
+        if (rewrittenQuery != query) {
-            return new HasParentQueryBuilder(type, rewrite, score, innerHit);
+            InnerHitBuilder rewrittenInnerHit = InnerHitBuilder.rewrite(innerHit, rewrittenQuery);
+            return new HasParentQueryBuilder(type, rewrittenQuery, score, rewrittenInnerHit);
         }
         return this;
     }

core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java

@@ -722,4 +722,16 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
         }
     }
 
+    static InnerHitBuilder rewrite(InnerHitBuilder original, QueryBuilder rewrittenQuery) {
+        if (original == null) {
+            return null;
+        }
+
+        InnerHitBuilder copy = new InnerHitBuilder(original);
+        copy.query = rewrittenQuery;
+        copy.parentChildType = original.parentChildType;
+        copy.nestedPath = original.nestedPath;
+        return copy;
+    }
+
 }

core/src/main/java/org/elasticsearch/index/query/NestedQueryBuilder.java

@@ -263,9 +263,10 @@ public class NestedQueryBuilder extends AbstractQueryBuilder<NestedQueryBuilder>
 
     @Override
     protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws IOException {
-        QueryBuilder rewrite = query.rewrite(queryRewriteContext);
+        QueryBuilder rewrittenQuery = query.rewrite(queryRewriteContext);
-        if (rewrite != query) {
+        if (rewrittenQuery != query) {
-            return new NestedQueryBuilder(path, rewrite, scoreMode, innerHitBuilder);
+            InnerHitBuilder rewrittenInnerHit = InnerHitBuilder.rewrite(innerHitBuilder, rewrittenQuery);
+            return new NestedQueryBuilder(path, rewrittenQuery, scoreMode, rewrittenInnerHit);
         }
         return this;
     }

core/src/test/java/org/elasticsearch/index/query/HasChildQueryBuilderTests.java

@@ -52,8 +52,6 @@ import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.sort.FieldSortBuilder;
 import org.elasticsearch.search.sort.SortOrder;
 import org.elasticsearch.test.AbstractQueryTestCase;
-import org.junit.Before;
-import org.junit.BeforeClass;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -74,6 +72,8 @@ public class HasChildQueryBuilderTests extends AbstractQueryTestCase<HasChildQue
 
     private static String similarity;
 
+    boolean requiresRewrite = false;
+
     @Override
     protected void initializeAdditionalMappings(MapperService mapperService) throws IOException {
         similarity = randomFrom("classic", "BM25");
@@ -105,8 +105,14 @@ public class HasChildQueryBuilderTests extends AbstractQueryTestCase<HasChildQue
     protected HasChildQueryBuilder doCreateTestQueryBuilder() {
         int min = randomIntBetween(0, Integer.MAX_VALUE / 2);
         int max = randomIntBetween(min, Integer.MAX_VALUE);
-        HasChildQueryBuilder hqb = new HasChildQueryBuilder(CHILD_TYPE,
+
-                RandomQueryBuilder.createQuery(random()),
+        QueryBuilder innerQueryBuilder = RandomQueryBuilder.createQuery(random());
+        if (randomBoolean()) {
+            requiresRewrite = true;
+            innerQueryBuilder = new WrapperQueryBuilder(innerQueryBuilder.toString());
+        }
+
+        HasChildQueryBuilder hqb = new HasChildQueryBuilder(CHILD_TYPE, innerQueryBuilder,
                 RandomPicks.randomFrom(random(), ScoreMode.values()));
         hqb.minMaxChildren(min, max);
         if (randomBoolean()) {
@@ -127,9 +133,11 @@ public class HasChildQueryBuilderTests extends AbstractQueryTestCase<HasChildQue
         assertEquals(queryBuilder.maxChildren(), lpq.getMaxChildren());
         assertEquals(queryBuilder.scoreMode(), lpq.getScoreMode()); // WTF is this why do we have two?
         if (queryBuilder.innerHit() != null) {
+            // have to rewrite again because the provided queryBuilder hasn't been rewritten (directly returned from
+            // doCreateTestQueryBuilder)
+            queryBuilder = (HasChildQueryBuilder) queryBuilder.rewrite(context);
             SearchContext searchContext = SearchContext.current();
             assertNotNull(searchContext);
-            if (query != null) {
             Map<String, InnerHitBuilder> innerHitBuilders = new HashMap<>();
             InnerHitBuilder.extractInnerHits(queryBuilder, innerHitBuilders);
             for (InnerHitBuilder builder : innerHitBuilders.values()) {
@@ -143,9 +151,6 @@ public class HasChildQueryBuilderTests extends AbstractQueryTestCase<HasChildQue
             assertEquals(innerHits.size(), queryBuilder.innerHit().getSize());
             assertEquals(innerHits.sort().sort.getSort().length, 1);
             assertEquals(innerHits.sort().sort.getSort()[0].getField(), STRING_FIELD_NAME_2);
-            } else {
-                assertThat(searchContext.innerHits().getInnerHits().size(), equalTo(0));
-            }
         }
     }
 
@@ -315,6 +320,17 @@ public class HasChildQueryBuilderTests extends AbstractQueryTestCase<HasChildQue
         }
     }
 
+    @Override
+    public void testMustRewrite() throws IOException {
+        try {
+            super.testMustRewrite();
+        } catch (UnsupportedOperationException e) {
+            if (requiresRewrite == false) {
+                throw e;
+            }
+        }
+    }
+
     public void testNonDefaultSimilarity() throws Exception {
         QueryShardContext shardContext = createShardContext();
         HasChildQueryBuilder hasChildQueryBuilder = QueryBuilders.hasChildQuery(CHILD_TYPE, new TermQueryBuilder("custom_string", "value"), ScoreMode.None);

core/src/test/java/org/elasticsearch/index/query/HasParentQueryBuilderTests.java

@@ -58,6 +58,8 @@ public class HasParentQueryBuilderTests extends AbstractQueryTestCase<HasParentQ
     protected static final String PARENT_TYPE = "parent";
     protected static final String CHILD_TYPE = "child";
 
+    boolean requiresRewrite = false;
+
     @Override
     protected void initializeAdditionalMappings(MapperService mapperService) throws IOException {
         mapperService.merge(PARENT_TYPE, new CompressedXContent(PutMappingRequest.buildFromSimplifiedDef(PARENT_TYPE,
@@ -88,8 +90,12 @@ public class HasParentQueryBuilderTests extends AbstractQueryTestCase<HasParentQ
      */
     @Override
     protected HasParentQueryBuilder doCreateTestQueryBuilder() {
-        HasParentQueryBuilder hqb = new HasParentQueryBuilder(PARENT_TYPE,
+        QueryBuilder innerQueryBuilder = RandomQueryBuilder.createQuery(random());
-                RandomQueryBuilder.createQuery(random()),randomBoolean());
+        if (randomBoolean()) {
+            requiresRewrite = true;
+            innerQueryBuilder = new WrapperQueryBuilder(innerQueryBuilder.toString());
+        }
+        HasParentQueryBuilder hqb = new HasParentQueryBuilder(PARENT_TYPE, innerQueryBuilder, randomBoolean());
         if (randomBoolean()) {
             hqb.innerHit(new InnerHitBuilder()
                     .setName(randomAsciiOfLengthBetween(1, 10))
@@ -107,9 +113,12 @@ public class HasParentQueryBuilderTests extends AbstractQueryTestCase<HasParentQ
         assertEquals(queryBuilder.score() ? ScoreMode.Max : ScoreMode.None, lpq.getScoreMode());
 
         if (queryBuilder.innerHit() != null) {
+            // have to rewrite again because the provided queryBuilder hasn't been rewritten (directly returned from
+            // doCreateTestQueryBuilder)
+            queryBuilder = (HasParentQueryBuilder) queryBuilder.rewrite(context);
+
             SearchContext searchContext = SearchContext.current();
             assertNotNull(searchContext);
-            if (query != null) {
             Map<String, InnerHitBuilder> innerHitBuilders = new HashMap<>();
             InnerHitBuilder.extractInnerHits(queryBuilder, innerHitBuilders);
             for (InnerHitBuilder builder : innerHitBuilders.values()) {
@@ -123,9 +132,6 @@ public class HasParentQueryBuilderTests extends AbstractQueryTestCase<HasParentQ
             assertEquals(innerHits.size(), queryBuilder.innerHit().getSize());
             assertEquals(innerHits.sort().sort.getSort().length, 1);
             assertEquals(innerHits.sort().sort.getSort()[0].getField(), STRING_FIELD_NAME_2);
-            } else {
-                assertThat(searchContext.innerHits().getInnerHits().size(), equalTo(0));
-            }
         }
     }
 
@@ -206,6 +212,17 @@ public class HasParentQueryBuilderTests extends AbstractQueryTestCase<HasParentQ
         }
     }
 
+    @Override
+    public void testMustRewrite() throws IOException {
+        try {
+            super.testMustRewrite();
+        } catch (UnsupportedOperationException e) {
+            if (requiresRewrite == false) {
+                throw e;
+            }
+        }
+    }
+
     public void testFromJson() throws IOException {
         String json =
                 "{\n" +

core/src/test/java/org/elasticsearch/index/query/NestedQueryBuilderTests.java

@@ -49,6 +49,8 @@ import static org.hamcrest.CoreMatchers.notNullValue;
 
 public class NestedQueryBuilderTests extends AbstractQueryTestCase<NestedQueryBuilder> {
 
+    boolean requiresRewrite = false;
+
     @Override
     protected void initializeAdditionalMappings(MapperService mapperService) throws IOException {
         mapperService.merge("nested_doc", new CompressedXContent(PutMappingRequest.buildFromSimplifiedDef("nested_doc",
@@ -68,7 +70,12 @@ public class NestedQueryBuilderTests extends AbstractQueryTestCase<NestedQueryBu
      */
     @Override
     protected NestedQueryBuilder doCreateTestQueryBuilder() {
-        NestedQueryBuilder nqb = new NestedQueryBuilder("nested1", RandomQueryBuilder.createQuery(random()),
+        QueryBuilder innerQueryBuilder = RandomQueryBuilder.createQuery(random());
+        if (randomBoolean()) {
+            requiresRewrite = true;
+            innerQueryBuilder = new WrapperQueryBuilder(innerQueryBuilder.toString());
+        }
+        NestedQueryBuilder nqb = new NestedQueryBuilder("nested1", innerQueryBuilder,
                 RandomPicks.randomFrom(random(), ScoreMode.values()));
         if (randomBoolean()) {
             nqb.innerHit(new InnerHitBuilder()
@@ -87,9 +94,12 @@ public class NestedQueryBuilderTests extends AbstractQueryTestCase<NestedQueryBu
         ToParentBlockJoinQuery parentBlockJoinQuery = (ToParentBlockJoinQuery) query;
         // TODO how to assert this?
         if (queryBuilder.innerHit() != null) {
+            // have to rewrite again because the provided queryBuilder hasn't been rewritten (directly returned from
+            // doCreateTestQueryBuilder)
+            queryBuilder = (NestedQueryBuilder) queryBuilder.rewrite(context);
+
             SearchContext searchContext = SearchContext.current();
             assertNotNull(searchContext);
-            if (query != null) {
             Map<String, InnerHitBuilder> innerHitBuilders = new HashMap<>();
             InnerHitBuilder.extractInnerHits(queryBuilder, innerHitBuilders);
             for (InnerHitBuilder builder : innerHitBuilders.values()) {
@@ -102,9 +112,6 @@ public class NestedQueryBuilderTests extends AbstractQueryTestCase<NestedQueryBu
             assertEquals(innerHits.size(), queryBuilder.innerHit().getSize());
             assertEquals(innerHits.sort().sort.getSort().length, 1);
             assertEquals(innerHits.sort().sort.getSort()[0].getField(), INT_FIELD_NAME);
-            } else {
-                assertThat(searchContext.innerHits().getInnerHits().size(), equalTo(0));
-            }
         }
     }
 
@@ -199,6 +206,17 @@ public class NestedQueryBuilderTests extends AbstractQueryTestCase<NestedQueryBu
         }
     }
 
+    @Override
+    public void testMustRewrite() throws IOException {
+        try {
+            super.testMustRewrite();
+        } catch (UnsupportedOperationException e) {
+            if (requiresRewrite == false) {
+                throw e;
+            }
+        }
+    }
+
     public void testIgnoreUnmapped() throws IOException {
         final NestedQueryBuilder queryBuilder = new NestedQueryBuilder("unmapped", new MatchAllQueryBuilder(), ScoreMode.None);
         queryBuilder.ignoreUnmapped(true);