SourceLoader.Leaf.source should return Source (#91352)

Most consumers of SourceLoader.Leaf.source need a Source object
rather than a plain BytesReference, so it makes sense to return that
directly.

server/src/main/java/org/elasticsearch/index/get/ShardGetService.java

@@ -241,7 +241,6 @@ public final class ShardGetService extends AbstractIndexShardComponent {
 
         Map<String, DocumentField> documentFields = null;
         Map<String, DocumentField> metadataFields = null;
-        BytesReference source;
         DocIdAndVersion docIdAndVersion = get.docIdAndVersion();
         SourceLoader loader = forceSyntheticSource
             ? new SourceLoader.Synthetic(mappingLookup.getMapping())
@@ -279,16 +278,16 @@ public final class ShardGetService extends AbstractIndexShardComponent {
                 }
             }
         }
-        source = loader.leaf(docIdAndVersion.reader, new int[] { docIdAndVersion.docId })
-            .source(leafStoredFieldLoader, docIdAndVersion.docId);
 
-        if (source != null) {
-            // apply request-level source filtering
-            if (fetchSourceContext.fetchSource() == false) {
-                source = null;
-            } else if (fetchSourceContext.hasFilter()) {
-                source = Source.fromBytes(source).filter(fetchSourceContext.filter()).internalSourceRef();
+        BytesReference sourceBytes = null;
+        if (mapperService.mappingLookup().isSourceEnabled() && fetchSourceContext.fetchSource()) {
+            Source source = loader.leaf(docIdAndVersion.reader, new int[] { docIdAndVersion.docId })
+                .source(leafStoredFieldLoader, docIdAndVersion.docId);
+
+            if (fetchSourceContext.hasFilter()) {
+                source = source.filter(fetchSourceContext.filter());
             }
+            sourceBytes = source.internalSourceRef();
         }
 
         return new GetResult(
@@ -298,7 +297,7 @@ public final class ShardGetService extends AbstractIndexShardComponent {
             get.docIdAndVersion().primaryTerm,
             get.version(),
             get.exists(),
-            source,
+            sourceBytes,
             documentFields,
             metadataFields
         );

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

@@ -11,6 +11,7 @@ package org.elasticsearch.index.mapper;
 import org.apache.lucene.index.LeafReader;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.index.fieldvisitor.LeafStoredFieldLoader;
+import org.elasticsearch.search.lookup.Source;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.json.JsonXContent;
 
@@ -51,14 +52,7 @@ public interface SourceLoader {
          * @param storedFields a loader for stored fields
          * @param docId the doc to load
          */
-        BytesReference source(LeafStoredFieldLoader storedFields, int docId) throws IOException;
-
-        Leaf EMPTY_OBJECT = (storedFields, docId) -> {
-            // TODO accept a requested xcontent type
-            try (XContentBuilder b = new XContentBuilder(JsonXContent.jsonXContent, new ByteArrayOutputStream())) {
-                return BytesReference.bytes(b.startObject().endObject());
-            }
-        };
+        Source source(LeafStoredFieldLoader storedFields, int docId) throws IOException;
     }
 
     /**
@@ -72,7 +66,7 @@ public interface SourceLoader {
 
         @Override
         public Leaf leaf(LeafReader reader, int[] docIdsInLeaf) {
-            return (storedFieldLoader, docId) -> storedFieldLoader.source();
+            return (storedFieldLoader, docId) -> Source.fromBytes(storedFieldLoader.source());
         }
 
         @Override
@@ -116,7 +110,7 @@ public interface SourceLoader {
             }
 
             @Override
-            public BytesReference source(LeafStoredFieldLoader storedFieldLoader, int docId) throws IOException {
+            public Source source(LeafStoredFieldLoader storedFieldLoader, int docId) throws IOException {
                 for (Map.Entry<String, List<Object>> e : storedFieldLoader.storedFields().entrySet()) {
                     SyntheticFieldLoader.StoredFieldLoader loader = storedFieldLoaders.get(e.getKey());
                     if (loader != null) {
@@ -133,7 +127,7 @@ public interface SourceLoader {
                     } else {
                         b.startObject().endObject();
                     }
-                    return BytesReference.bytes(b);
+                    return Source.fromBytes(BytesReference.bytes(b), b.contentType());
                 }
             }
         }

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

@@ -22,6 +22,7 @@ import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.text.Text;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.common.xcontent.XContentHelper;
+import org.elasticsearch.common.xcontent.support.XContentMapValues;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.RestApiVersion;
 import org.elasticsearch.index.mapper.IgnoredFieldMapper;
@@ -31,6 +32,7 @@ import org.elasticsearch.index.seqno.SequenceNumbers;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.search.fetch.subphase.LookupField;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
+import org.elasticsearch.search.lookup.Source;
 import org.elasticsearch.search.lookup.SourceLookup;
 import org.elasticsearch.transport.RemoteClusterAware;
 import org.elasticsearch.xcontent.ConstructingObjectParser;
@@ -42,6 +44,7 @@ import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xcontent.XContentParser.Token;
+import org.elasticsearch.xcontent.XContentType;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -1071,6 +1074,51 @@ public final class SearchHit implements Writeable, ToXContentObject, Iterable<Do
             return child;
         }
 
+        /**
+         * Extracts the part of the root source that applies to this particular NestedIdentity, while
+         * preserving the enclosing path structure.
+         *
+         * For a root document that looks like this:
+         * { "children" :
+         *    [
+         *      { "grandchildren" : [ { "field" : "value1" }, { "field" : "value2" } ] },
+         *      { "grandchildren" : [ { "field" : "value3" }, { "field" : "value4" } ] }
+         *   ]
+         * }
+         *
+         * Extracting the NestedIdentity of the first child and second grandchild results in a source that looks like this:
+         * { "children" : { "grandchildren" : { "field" : "value2" } } }
+         *
+         * If the relevant child source object does not exist in the root, then we return {@link Source#empty(XContentType)}
+         */
+        @SuppressWarnings("unchecked")
+        public Source extractSource(Source root) {
+            // Isolate the nested json array object that matches with nested hit and wrap it back into the same json
+            // structure with the nested json array object being the actual content. The latter is important, so that
+            // features like source filtering and highlighting work consistent regardless of whether the field points
+            // to a json object array for consistency reasons on how we refer to fields
+            Map<String, Object> rootSourceAsMap = root.source();
+            Map<String, Object> nestedSourceAsMap = new HashMap<>();
+            Map<String, Object> current = nestedSourceAsMap;
+            for (SearchHit.NestedIdentity nested = this; nested != null; nested = nested.getChild()) {
+                String nestedPath = nested.getField().string();
+                current.put(nestedPath, new HashMap<>());
+                List<Map<?, ?>> nestedParsedSource = XContentMapValues.extractNestedSources(nestedPath, rootSourceAsMap);
+                if (nestedParsedSource == null) {
+                    return Source.empty(root.sourceContentType());
+                }
+                rootSourceAsMap = (Map<String, Object>) nestedParsedSource.get(nested.getOffset());
+                if (nested.getChild() == null) {
+                    current.put(nestedPath, rootSourceAsMap);
+                } else {
+                    Map<String, Object> next = new HashMap<>();
+                    current.put(nestedPath, next);
+                    current = next;
+                }
+            }
+            return Source.fromMap(nestedSourceAsMap, root.sourceContentType());
+        }
+
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             out.writeOptionalText(field);

server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java

@@ -12,9 +12,6 @@ import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.TotalHits;
-import org.elasticsearch.common.xcontent.XContentHelper;
-import org.elasticsearch.common.xcontent.support.XContentMapValues;
-import org.elasticsearch.core.Tuple;
 import org.elasticsearch.index.fieldvisitor.LeafStoredFieldLoader;
 import org.elasticsearch.index.fieldvisitor.StoredFieldLoader;
 import org.elasticsearch.index.mapper.SourceLoader;
@@ -39,7 +36,6 @@ import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Supplier;
@@ -243,7 +239,7 @@ public class FetchPhase {
             if (requiresSource) {
                 try {
                     profiler.startLoadingSource();
-                    source = Source.fromBytes(sourceLoader.source(leafStoredFieldLoader, subDocId));
+                    source = sourceLoader.source(leafStoredFieldLoader, subDocId);
                     SourceLookup scriptSourceLookup = context.getSearchExecutionContext().lookup().source();
                     scriptSourceLookup.setSegmentAndDocument(subReaderContext, subDocId);
                     scriptSourceLookup.setSourceProvider(new SourceLookup.BytesSourceProvider(source.internalSourceRef()));
@@ -278,7 +274,6 @@ public class FetchPhase {
      *     setting it on {@link HitContext#source()}. This allows fetch subphases that
      *     use the hit context to access the preloaded source.
      */
-    @SuppressWarnings("unchecked")
     private static HitContext prepareNestedHitContext(
         SearchContext context,
         boolean requiresSource,
@@ -290,16 +285,13 @@ public class FetchPhase {
     ) throws IOException {
 
         String rootId;
-        Map<String, Object> rootSourceAsMap = null;
-        XContentType rootSourceContentType = null;
+        Source rootSource = Source.empty(XContentType.JSON);
 
         if (context instanceof InnerHitsContext.InnerHitSubContext innerHitsContext) {
             rootId = innerHitsContext.getRootId();
 
             if (requiresSource) {
-                Source rootLookup = innerHitsContext.getRootLookup();
-                rootSourceAsMap = rootLookup.source();
-                rootSourceContentType = rootLookup.sourceContentType();
+                rootSource = innerHitsContext.getRootLookup();
             }
         } else {
             StoredFieldLoader rootLoader = profiler.storedFields(StoredFieldLoader.create(requiresSource, Collections.emptySet()));
@@ -309,11 +301,7 @@ public class FetchPhase {
 
             if (requiresSource) {
                 if (leafRootLoader.source() != null) {
-                    Tuple<XContentType, Map<String, Object>> tuple = XContentHelper.convertToMap(leafRootLoader.source(), false);
-                    rootSourceAsMap = tuple.v2();
-                    rootSourceContentType = tuple.v1();
-                } else {
-                    rootSourceAsMap = Collections.emptyMap();
+                    rootSource = Source.fromBytes(leafRootLoader.source());
                 }
             }
         }
@@ -321,41 +309,11 @@ public class FetchPhase {
         childFieldLoader.advanceTo(nestedInfo.doc());
 
         SearchHit.NestedIdentity nestedIdentity = nestedInfo.nestedIdentity();
+        assert nestedIdentity != null;
+        Source nestedSource = nestedIdentity.extractSource(rootSource);
 
         SearchHit hit = new SearchHit(topDocId, rootId, nestedIdentity);
-
-        if (rootSourceAsMap != null && rootSourceAsMap.isEmpty() == false) {
-            // Isolate the nested json array object that matches with nested hit and wrap it back into the same json
-            // structure with the nested json array object being the actual content. The latter is important, so that
-            // features like source filtering and highlighting work consistent regardless of whether the field points
-            // to a json object array for consistency reasons on how we refer to fields
-            Map<String, Object> nestedSourceAsMap = new HashMap<>();
-            Map<String, Object> current = nestedSourceAsMap;
-            for (SearchHit.NestedIdentity nested = nestedIdentity; nested != null; nested = nested.getChild()) {
-                String nestedPath = nested.getField().string();
-                current.put(nestedPath, new HashMap<>());
-                List<Map<?, ?>> nestedParsedSource = XContentMapValues.extractNestedSources(nestedPath, rootSourceAsMap);
-                if (nestedParsedSource == null) {
-                    throw new IllegalStateException("Couldn't find nested source for path " + nestedPath);
-                }
-                rootSourceAsMap = (Map<String, Object>) nestedParsedSource.get(nested.getOffset());
-                if (nested.getChild() == null) {
-                    current.put(nestedPath, rootSourceAsMap);
-                } else {
-                    Map<String, Object> next = new HashMap<>();
-                    current.put(nestedPath, next);
-                    current = next;
-                }
-            }
-            return new HitContext(
-                hit,
-                subReaderContext,
-                nestedInfo.doc(),
-                childFieldLoader.storedFields(),
-                Source.fromMap(nestedSourceAsMap, rootSourceContentType)
-            );
-        }
-        return new HitContext(hit, subReaderContext, nestedInfo.doc(), childFieldLoader.storedFields(), Source.empty(null));
+        return new HitContext(hit, subReaderContext, nestedInfo.doc(), childFieldLoader.storedFields(), nestedSource);
     }
 
     interface Profiler {

test/framework/src/main/java/org/elasticsearch/index/mapper/MapperServiceTestCase.java

@@ -62,6 +62,7 @@ import org.elasticsearch.search.aggregations.support.AggregationContext;
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.search.internal.SubSearchContext;
 import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.search.lookup.Source;
 import org.elasticsearch.search.sort.BucketedSort;
 import org.elasticsearch.search.sort.BucketedSort.ExtraData;
 import org.elasticsearch.search.sort.SortAndFormats;
@@ -715,8 +716,8 @@ public abstract class MapperServiceTestCase extends ESTestCase {
         SourceLoader loader = mapper.sourceMapper().newSourceLoader(mapper.mapping());
         LeafReader leafReader = getOnlyLeafReader(reader);
         SourceLoader.Leaf leafLoader = loader.leaf(leafReader, new int[] { docId });
-        BytesReference syntheticSourceBytes = leafLoader.source(syntheticSourceStoredFieldLoader(mapper, leafReader, loader), docId);
-        return syntheticSourceBytes.utf8ToString();
+        Source synthetic = leafLoader.source(syntheticSourceStoredFieldLoader(mapper, leafReader, loader), docId);
+        return synthetic.internalSourceRef().utf8ToString();
     }
 
     protected static LeafStoredFieldLoader syntheticSourceStoredFieldLoader(

test/framework/src/main/java/org/elasticsearch/index/mapper/MapperTestCase.java

@@ -1094,7 +1094,8 @@ public abstract class MapperTestCase extends MapperServiceTestCase {
                     LeafStoredFieldLoader storedLeaf = storedFieldLoader.getLoader(leaf, docIds);
                     for (int docId : docIds) {
                         storedLeaf.advanceTo(docId);
-                        assertThat("doc " + docId, sourceLoaderLeaf.source(storedLeaf, docId).utf8ToString(), equalTo(expected[i++]));
+                        String source = sourceLoaderLeaf.source(storedLeaf, docId).internalSourceRef().utf8ToString();
+                        assertThat("doc " + docId, source, equalTo(expected[i++]));
                     }
                 }
             }