Add the ability to define custom inner hit sub context builder (#24676)

This commit moves the handling of nested and parent/child inner hits to specialized classes that can be defined outside of ES core.
InnerHitBuilderContext is now used by the parent query (nested or hasChild, ...) to build the sub context from the InnerHitBuilder definition.
BWC is also ensured so that nodes in previous versions can still send/receive inner hits to/from this version.

Relates #20257

1 files changed, 100 insertions, 280 deletions

diff --git a/core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java b/core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java
index d21d3d6670..4419d6a93a 100644
--- a/core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java
+++ b/core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java
@@ -28,31 +28,21 @@ import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.xcontent.ObjectParser;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.index.mapper.DocumentMapper;
-import org.elasticsearch.index.mapper.ObjectMapper;
 import org.elasticsearch.script.Script;
-import org.elasticsearch.script.ScriptContext;
-import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.builder.SearchSourceBuilder.ScriptField;
 import org.elasticsearch.search.fetch.StoredFieldsContext;
-import org.elasticsearch.search.fetch.subphase.DocValueFieldsContext;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
-import org.elasticsearch.search.fetch.subphase.InnerHitsContext;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
-import org.elasticsearch.search.internal.SearchContext;
-import org.elasticsearch.search.sort.SortAndFormats;
 import org.elasticsearch.search.sort.SortBuilder;
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.Comparator;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
-import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.Set;
 
 import static org.elasticsearch.common.xcontent.XContentParser.Token.END_OBJECT;
@@ -61,7 +51,6 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
 
     public static final ParseField NAME_FIELD = new ParseField("name");
     public static final ParseField IGNORE_UNMAPPED = new ParseField("ignore_unmapped");
-    public static final ParseField INNER_HITS_FIELD = new ParseField("inner_hits");
     public static final QueryBuilder DEFAULT_INNER_HIT_QUERY = new MatchAllQueryBuilder();
 
     private static final ObjectParser<InnerHitBuilder, QueryParseContext> PARSER = new ObjectParser<>("inner_hits", InnerHitBuilder::new);
@@ -104,35 +93,9 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
         }, SearchSourceBuilder._SOURCE_FIELD, ObjectParser.ValueType.OBJECT_ARRAY_BOOLEAN_OR_STRING);
         PARSER.declareObject(InnerHitBuilder::setHighlightBuilder, (p, c) -> HighlightBuilder.fromXContent(c),
                 SearchSourceBuilder.HIGHLIGHT_FIELD);
-        PARSER.declareObject(InnerHitBuilder::setChildInnerHits, (p, c) -> {
-            try {
-                Map<String, InnerHitBuilder> innerHitBuilders = new HashMap<>();
-                String innerHitName = null;
-                for (XContentParser.Token token = p.nextToken(); token != XContentParser.Token.END_OBJECT; token = p.nextToken()) {
-                    switch (token) {
-                        case START_OBJECT:
-                            InnerHitBuilder innerHitBuilder = InnerHitBuilder.fromXContent(c);
-                            innerHitBuilder.setName(innerHitName);
-                            innerHitBuilders.put(innerHitName, innerHitBuilder);
-                            break;
-                        case FIELD_NAME:
-                            innerHitName = p.currentName();
-                            break;
-                        default:
-                            throw new ParsingException(p.getTokenLocation(), "Expected [" + XContentParser.Token.START_OBJECT + "] in ["
-                                    + p.currentName() + "] but found [" + token + "]", p.getTokenLocation());
-                    }
-                }
-                return innerHitBuilders;
-            } catch (IOException e) {
-                throw new ParsingException(p.getTokenLocation(), "Could not parse inner query definition", e);
-            }
-        }, INNER_HITS_FIELD);
     }
 
     private String name;
-    private String nestedPath;
-    private String parentChildType;
     private boolean ignoreUnmapped;
 
     private int from;
@@ -148,71 +111,25 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
     private Set<ScriptField> scriptFields;
     private HighlightBuilder highlightBuilder;
     private FetchSourceContext fetchSourceContext;
-    private Map<String, InnerHitBuilder> childInnerHits;
 
     public InnerHitBuilder() {
+        this.name = null;
     }
 
-    private InnerHitBuilder(InnerHitBuilder other) {
-        name = other.name;
-        this.ignoreUnmapped = other.ignoreUnmapped;
-        from = other.from;
-        size = other.size;
-        explain = other.explain;
-        version = other.version;
-        trackScores = other.trackScores;
-        if (other.storedFieldsContext != null) {
-            storedFieldsContext = new StoredFieldsContext(other.storedFieldsContext);
-        }
-        if (other.docValueFields != null) {
-            docValueFields = new ArrayList<> (other.docValueFields);
-        }
-        if (other.scriptFields != null) {
-            scriptFields = new HashSet<> (other.scriptFields);
-        }
-        if (other.fetchSourceContext != null) {
-            fetchSourceContext = new FetchSourceContext(
-                    other.fetchSourceContext.fetchSource(), other.fetchSourceContext.includes(), other.fetchSourceContext.excludes()
-            );
-        }
-        if (other.sorts != null) {
-            sorts = new ArrayList<>(other.sorts);
-        }
-        highlightBuilder = other.highlightBuilder;
-        if (other.childInnerHits != null) {
-            childInnerHits = new HashMap<>(other.childInnerHits);
-        }
-    }
-
-
-    InnerHitBuilder(InnerHitBuilder other, String nestedPath, QueryBuilder query, boolean ignoreUnmapped) {
-        this(other);
-        this.query = query;
-        this.nestedPath = nestedPath;
-        this.ignoreUnmapped = ignoreUnmapped;
-        if (name == null) {
-            this.name = nestedPath;
-        }
+    public InnerHitBuilder(String name) {
+        this.name = name;
     }
 
-    // NORELEASE Do not use this ctr, it is public for hasChild and hasParent query but this is temporary
-    public InnerHitBuilder(InnerHitBuilder other, QueryBuilder query, String parentChildType, boolean ignoreUnmapped) {
-        this(other);
-        this.query = query;
-        this.parentChildType = parentChildType;
-        this.ignoreUnmapped = ignoreUnmapped;
-        if (name == null) {
-            this.name = parentChildType;
-        }
-    }
 
     /**
      * Read from a stream.
      */
     public InnerHitBuilder(StreamInput in) throws IOException {
         name = in.readOptionalString();
-        nestedPath = in.readOptionalString();
-        parentChildType = in.readOptionalString();
+        if (in.getVersion().before(Version.V_6_0_0_alpha2_UNRELEASED)) {
+            in.readOptionalString();
+            in.readOptionalString();
+        }
         if (in.getVersion().onOrAfter(Version.V_5_2_0)) {
             ignoreUnmapped = in.readBoolean();
         }
@@ -239,21 +156,94 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
             }
         }
         highlightBuilder = in.readOptionalWriteable(HighlightBuilder::new);
-        query = in.readNamedWriteable(QueryBuilder.class);
-        if (in.readBoolean()) {
-            int size = in.readVInt();
-            childInnerHits = new HashMap<>(size);
-            for (int i = 0; i < size; i++) {
-                childInnerHits.put(in.readString(), new InnerHitBuilder(in));
-            }
+        if (in.getVersion().before(Version.V_6_0_0_alpha2_UNRELEASED)) {
+            /**
+             * this is needed for BWC with nodes pre 5.5
+             */
+            in.readNamedWriteable(QueryBuilder.class);
+            boolean hasChildren = in.readBoolean();
+            assert hasChildren == false;
         }
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
+        if (out.getVersion().before(Version.V_6_0_0_alpha2_UNRELEASED)) {
+            throw new IOException("Invalid output version, must >= " + Version.V_6_0_0_alpha2_UNRELEASED.toString());
+        }
         out.writeOptionalString(name);
-        out.writeOptionalString(nestedPath);
-        out.writeOptionalString(parentChildType);
+        out.writeBoolean(ignoreUnmapped);
+        out.writeVInt(from);
+        out.writeVInt(size);
+        out.writeBoolean(explain);
+        out.writeBoolean(version);
+        out.writeBoolean(trackScores);
+        out.writeOptionalWriteable(storedFieldsContext);
+        out.writeGenericValue(docValueFields);
+        boolean hasScriptFields = scriptFields != null;
+        out.writeBoolean(hasScriptFields);
+        if (hasScriptFields) {
+            out.writeVInt(scriptFields.size());
+            Iterator<ScriptField> iterator = scriptFields.stream()
+                    .sorted(Comparator.comparing(ScriptField::fieldName)).iterator();
+            while (iterator.hasNext()) {
+                iterator.next().writeTo(out);
+            }
+        }
+        out.writeOptionalWriteable(fetchSourceContext);
+        boolean hasSorts = sorts != null;
+        out.writeBoolean(hasSorts);
+        if (hasSorts) {
+            out.writeVInt(sorts.size());
+            for (SortBuilder<?> sort : sorts) {
+                out.writeNamedWriteable(sort);
+            }
+        }
+        out.writeOptionalWriteable(highlightBuilder);
+    }
+
+    /**
+     * BWC serialization for nested {@link InnerHitBuilder}.
+     * Should only be used to send nested inner hits to nodes pre 5.5.
+     */
+    protected void writeToNestedBWC(StreamOutput out, QueryBuilder query, String nestedPath) throws IOException {
+        assert out.getVersion().before(Version.V_6_0_0_alpha2_UNRELEASED) :
+            "invalid output version, must be < " + Version.V_6_0_0_alpha2_UNRELEASED.toString();
+        writeToBWC(out, query, nestedPath, null);
+    }
+
+    /**
+     * BWC serialization for collapsing {@link InnerHitBuilder}.
+     * Should only be used to send collapsing inner hits to nodes pre 5.5.
+     */
+    public void writeToCollapseBWC(StreamOutput out) throws IOException {
+        assert out.getVersion().before(Version.V_6_0_0_alpha2_UNRELEASED) :
+            "invalid output version, must be < " + Version.V_6_0_0_alpha2_UNRELEASED.toString();
+        writeToBWC(out, new MatchAllQueryBuilder(), null, null);
+    }
+
+    /**
+     * BWC serialization for parent/child {@link InnerHitBuilder}.
+     * Should only be used to send hasParent or hasChild inner hits to nodes pre 5.5.
+     */
+    public void writeToParentChildBWC(StreamOutput out, QueryBuilder query, String parentChildPath) throws IOException {
+        assert(out.getVersion().before(Version.V_6_0_0_alpha2_UNRELEASED)) :
+            "invalid output version, must be < " + Version.V_6_0_0_alpha2_UNRELEASED.toString();
+        writeToBWC(out, query, null, parentChildPath);
+    }
+
+    private void writeToBWC(StreamOutput out,
+                            QueryBuilder query,
+                            String nestedPath,
+                            String parentChildPath) throws IOException {
+        out.writeOptionalString(name);
+        if (nestedPath != null) {
+            out.writeOptionalString(nestedPath);
+            out.writeOptionalString(null);
+        } else {
+            out.writeOptionalString(null);
+            out.writeOptionalString(parentChildPath);
+        }
         if (out.getVersion().onOrAfter(Version.V_5_2_0)) {
             out.writeBoolean(ignoreUnmapped);
         }
@@ -269,7 +259,7 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
         if (hasScriptFields) {
             out.writeVInt(scriptFields.size());
             Iterator<ScriptField> iterator = scriptFields.stream()
-                    .sorted((a, b) -> a.fieldName().compareTo(b.fieldName())).iterator();
+                .sorted(Comparator.comparing(ScriptField::fieldName)).iterator();
             while (iterator.hasNext()) {
                 iterator.next().writeTo(out);
             }
@@ -285,18 +275,7 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
         }
         out.writeOptionalWriteable(highlightBuilder);
         out.writeNamedWriteable(query);
-        boolean hasChildInnerHits = childInnerHits != null;
-        out.writeBoolean(hasChildInnerHits);
-        if (hasChildInnerHits) {
-            out.writeVInt(childInnerHits.size());
-            Iterator<Map.Entry<String, InnerHitBuilder>> iterator = childInnerHits.entrySet().stream()
-                    .sorted((a, b) -> a.getKey().compareTo(b.getKey())).iterator();
-            while (iterator.hasNext()) {
-                Map.Entry<String, InnerHitBuilder> entry = iterator.next();
-                out.writeString(entry.getKey());
-                entry.getValue().writeTo(out);
-            }
-        }
+        out.writeBoolean(false);
     }
 
     public String getName() {
@@ -308,6 +287,11 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
         return this;
     }
 
+    public InnerHitBuilder setIgnoreUnmapped(boolean value) {
+        this.ignoreUnmapped = value;
+        return this;
+    }
+
     /**
      * Whether to include inner hits in the search response hits if required mappings is missing
      */
@@ -525,136 +509,6 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
         return query;
     }
 
-    void setChildInnerHits(Map<String, InnerHitBuilder> childInnerHits) {
-        this.childInnerHits = childInnerHits;
-    }
-
-    String getParentChildType() {
-        return parentChildType;
-    }
-
-    String getNestedPath() {
-        return nestedPath;
-    }
-
-    void addChildInnerHit(InnerHitBuilder innerHitBuilder) {
-        if (childInnerHits == null) {
-            childInnerHits = new HashMap<>();
-        }
-        this.childInnerHits.put(innerHitBuilder.getName(), innerHitBuilder);
-    }
-
-    public InnerHitsContext.BaseInnerHits build(SearchContext parentSearchContext,
-                                                InnerHitsContext innerHitsContext) throws IOException {
-        QueryShardContext queryShardContext = parentSearchContext.getQueryShardContext();
-        if (nestedPath != null) {
-            ObjectMapper nestedObjectMapper = queryShardContext.getObjectMapper(nestedPath);
-            if (nestedObjectMapper == null) {
-                if (ignoreUnmapped == false) {
-                    throw new IllegalStateException("[" + query.getName() + "] no mapping found for type [" + nestedPath + "]");
-                } else {
-                    return null;
-                }
-            }
-
-            ObjectMapper parentObjectMapper = queryShardContext.nestedScope().nextLevel(nestedObjectMapper);
-            InnerHitsContext.NestedInnerHits nestedInnerHits = new InnerHitsContext.NestedInnerHits(
-                    name, parentSearchContext, parentObjectMapper, nestedObjectMapper
-            );
-            setupInnerHitsContext(queryShardContext, nestedInnerHits);
-            if (childInnerHits != null) {
-                buildChildInnerHits(parentSearchContext, nestedInnerHits);
-            }
-            queryShardContext.nestedScope().previousLevel();
-            innerHitsContext.addInnerHitDefinition(nestedInnerHits);
-            return nestedInnerHits;
-        } else if (parentChildType != null) {
-            DocumentMapper documentMapper = queryShardContext.documentMapper(parentChildType);
-            if (documentMapper == null) {
-                if (ignoreUnmapped == false) {
-                    throw new IllegalStateException("[" + query.getName() + "] no mapping found for type [" + parentChildType + "]");
-                } else {
-                    return null;
-                }
-            }
-
-            InnerHitsContext.ParentChildInnerHits parentChildInnerHits = new InnerHitsContext.ParentChildInnerHits(
-                    name, parentSearchContext, queryShardContext.getMapperService(), documentMapper
-            );
-            setupInnerHitsContext(queryShardContext, parentChildInnerHits);
-            if (childInnerHits != null) {
-                buildChildInnerHits(parentSearchContext, parentChildInnerHits);
-            }
-            innerHitsContext.addInnerHitDefinition( parentChildInnerHits);
-            return parentChildInnerHits;
-        } else {
-            throw new IllegalStateException("Neither a nested or parent/child inner hit");
-        }
-    }
-
-    private void buildChildInnerHits(SearchContext parentSearchContext, InnerHitsContext.BaseInnerHits innerHits) throws IOException {
-        Map<String, InnerHitsContext.BaseInnerHits> childInnerHits = new HashMap<>();
-        for (Map.Entry<String, InnerHitBuilder> entry : this.childInnerHits.entrySet()) {
-            InnerHitsContext.BaseInnerHits childInnerHit = entry.getValue().build(
-                    parentSearchContext, new InnerHitsContext()
-            );
-            if (childInnerHit != null) {
-                childInnerHits.put(entry.getKey(), childInnerHit);
-            }
-        }
-        innerHits.setChildInnerHits(childInnerHits);
-    }
-
-    private void setupInnerHitsContext(QueryShardContext context, InnerHitsContext.BaseInnerHits innerHitsContext) throws IOException {
-        innerHitsContext.from(from);
-        innerHitsContext.size(size);
-        innerHitsContext.explain(explain);
-        innerHitsContext.version(version);
-        innerHitsContext.trackScores(trackScores);
-        if (storedFieldsContext != null) {
-            innerHitsContext.storedFieldsContext(storedFieldsContext);
-        }
-        if (docValueFields != null) {
-            innerHitsContext.docValueFieldsContext(new DocValueFieldsContext(docValueFields));
-        }
-        if (scriptFields != null) {
-            for (ScriptField field : scriptFields) {
-                SearchScript searchScript = innerHitsContext.getQueryShardContext().getSearchScript(field.script(),
-                    ScriptContext.SEARCH);
-                innerHitsContext.scriptFields().add(new org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField(
-                        field.fieldName(), searchScript, field.ignoreFailure()));
-            }
-        }
-        if (fetchSourceContext != null) {
-            innerHitsContext.fetchSourceContext(fetchSourceContext);
-        }
-        if (sorts != null) {
-            Optional<SortAndFormats> optionalSort = SortBuilder.buildSort(sorts, context);
-            if (optionalSort.isPresent()) {
-                innerHitsContext.sort(optionalSort.get());
-            }
-        }
-        if (highlightBuilder != null) {
-            innerHitsContext.highlight(highlightBuilder.build(context));
-        }
-        ParsedQuery parsedQuery = new ParsedQuery(query.toQuery(context), context.copyNamedQueries());
-        innerHitsContext.parsedQuery(parsedQuery);
-    }
-
-    public void inlineInnerHits(Map<String, InnerHitBuilder> innerHits) {
-        InnerHitBuilder copy = new InnerHitBuilder(this);
-        copy.parentChildType = this.parentChildType;
-        copy.nestedPath = this.nestedPath;
-        copy.query = this.query;
-        innerHits.put(copy.getName(), copy);
-
-        Map<String, InnerHitBuilder> childInnerHits = new HashMap<>();
-        extractInnerHits(query, childInnerHits);
-        if (childInnerHits.size() > 0) {
-            copy.setChildInnerHits(childInnerHits);
-        }
-    }
-
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
@@ -697,13 +551,6 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
         if (highlightBuilder != null) {
             builder.field(SearchSourceBuilder.HIGHLIGHT_FIELD.getPreferredName(), highlightBuilder, params);
         }
-        if (childInnerHits != null) {
-            builder.startObject(INNER_HITS_FIELD.getPreferredName());
-            for (Map.Entry<String, InnerHitBuilder> entry : childInnerHits.entrySet()) {
-                builder.field(entry.getKey(), entry.getValue(), params);
-            }
-            builder.endObject();
-        }
         builder.endObject();
         return builder;
     }
@@ -715,8 +562,6 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
 
         InnerHitBuilder that = (InnerHitBuilder) o;
         return Objects.equals(name, that.name) &&
-                Objects.equals(nestedPath, that.nestedPath) &&
-                Objects.equals(parentChildType, that.parentChildType) &&
                 Objects.equals(ignoreUnmapped, that.ignoreUnmapped) &&
                 Objects.equals(from, that.from) &&
                 Objects.equals(size, that.size) &&
@@ -728,41 +573,16 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
                 Objects.equals(scriptFields, that.scriptFields) &&
                 Objects.equals(fetchSourceContext, that.fetchSourceContext) &&
                 Objects.equals(sorts, that.sorts) &&
-                Objects.equals(highlightBuilder, that.highlightBuilder) &&
-                Objects.equals(query, that.query) &&
-                Objects.equals(childInnerHits, that.childInnerHits);
+                Objects.equals(highlightBuilder, that.highlightBuilder);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(name, nestedPath, parentChildType, ignoreUnmapped, from, size, explain, version, trackScores,
-                storedFieldsContext, docValueFields, scriptFields, fetchSourceContext, sorts, highlightBuilder, query, childInnerHits);
+        return Objects.hash(name, ignoreUnmapped, from, size, explain, version, trackScores,
+                storedFieldsContext, docValueFields, scriptFields, fetchSourceContext, sorts, highlightBuilder);
     }
 
     public static InnerHitBuilder fromXContent(QueryParseContext context) throws IOException {
         return PARSER.parse(context.parser(), new InnerHitBuilder(), context);
     }
-
-    public static void extractInnerHits(QueryBuilder query, Map<String, InnerHitBuilder> innerHitBuilders) {
-        if (query instanceof AbstractQueryBuilder) {
-            ((AbstractQueryBuilder) query).extractInnerHitBuilders(innerHitBuilders);
-        } else {
-            throw new IllegalStateException("provided query builder [" + query.getClass() +
-                    "] class should inherit from AbstractQueryBuilder, but it doesn't");
-        }
-    }
-
-    // TODO public for hasParent and hasChild query
-    public 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;
-    }
-
 }