Always build SourceProvider from the mapping lookup

Author: jimczi

benchmarks/src/main/java/org/elasticsearch/benchmark/script/ScriptScoreBenchmark.java

@@ -30,8 +30,10 @@
 import org.elasticsearch.index.fielddata.IndexFieldDataCache;
 import org.elasticsearch.index.fielddata.IndexNumericFieldData;
 import org.elasticsearch.index.mapper.MappedFieldType;
+import org.elasticsearch.index.mapper.MappingLookup;
 import org.elasticsearch.index.mapper.NumberFieldMapper.NumberFieldType;
 import org.elasticsearch.index.mapper.NumberFieldMapper.NumberType;
+import org.elasticsearch.index.mapper.SourceFieldMetrics;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
 import org.elasticsearch.indices.breaker.NoneCircuitBreakerService;
 import org.elasticsearch.plugins.PluginsService;
@@ -89,7 +91,7 @@ public class ScriptScoreBenchmark {
     private final SearchLookup lookup = new SearchLookup(
         fieldTypes::get,
         (mft, lookup, fdo) -> mft.fielddataBuilder(FieldDataContext.noRuntimeFields("benchmark")).build(fieldDataCache, breakerService),
-        SourceProvider.fromStoredFields()
+        SourceProvider.fromLookup(MappingLookup.EMPTY, SourceFieldMetrics.NOOP)
     );
 
     @Param({ "expression", "metal", "painless_cast", "painless_def" })

plugins/mapper-murmur3/src/test/java/org/elasticsearch/index/mapper/murmur3/Murmur3FieldMapperTests.java

@@ -107,7 +107,8 @@ protected void assertFetch(MapperService mapperService, String field, Object val
         ValueFetcher nativeFetcher = ft.valueFetcher(searchExecutionContext, format);
         ParsedDocument doc = mapperService.documentMapper().parse(source);
         withLuceneIndex(mapperService, iw -> iw.addDocuments(doc.docs()), ir -> {
-            Source s = SourceProvider.fromStoredFields().getSource(ir.leaves().get(0), 0);
+            Source s = SourceProvider.fromLookup(mapperService.mappingLookup(), mapperService.getMapperMetrics().sourceFieldMetrics())
+                .getSource(ir.leaves().get(0), 0);
             docValueFetcher.setNextReader(ir.leaves().get(0));
             nativeFetcher.setNextReader(ir.leaves().get(0));
             List<Object> fromDocValues = docValueFetcher.fetchValues(s, 0, new ArrayList<>());

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

@@ -493,9 +493,7 @@ public boolean containsBrokenAnalysis(String field) {
      */
     public SearchLookup lookup() {
         if (this.lookup == null) {
-            SourceProvider sourceProvider = isSourceSynthetic()
-                ? SourceProvider.fromSyntheticSource(mappingLookup.getMapping(), mapperMetrics.sourceFieldMetrics())
-                : SourceProvider.fromStoredFields();
+            SourceProvider sourceProvider = SourceProvider.fromLookup(mappingLookup, mapperMetrics.sourceFieldMetrics());
             setLookupProviders(sourceProvider, LeafFieldLookupProvider.fromStoredFields());
         }
         return this.lookup;

server/src/main/java/org/elasticsearch/search/lookup/ConcurrentSegmentSourceProvider.java

@@ -0,0 +1,73 @@
+/*
+ * 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.search.lookup;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
+import org.elasticsearch.index.fieldvisitor.LeafStoredFieldLoader;
+import org.elasticsearch.index.fieldvisitor.StoredFieldLoader;
+import org.elasticsearch.index.mapper.SourceFieldMapper;
+import org.elasticsearch.index.mapper.SourceLoader;
+
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ * A {@link SourceProvider} that loads _source from a concurrent search.
+ *
+ * @warning: This is written under the assumption that individual segments are accessed by a single
+ * thread, even if separate segments may be searched concurrently. If we ever implement
+ * within-segment concurrency this will have to work entirely differently.
+ * **/
+class ConcurrentSegmentSourceProvider implements SourceProvider {
+    private final SourceLoader sourceLoader;
+    private final StoredFieldLoader storedFieldLoader;
+    private final Map<Object, Leaf> leaves = ConcurrentCollections.newConcurrentMap();
+
+    ConcurrentSegmentSourceProvider(SourceLoader loader) {
+        this.sourceLoader = loader;
+        var fields = sourceLoader.requiredStoredFields();
+        boolean loadSource = fields.contains(SourceFieldMapper.NAME);
+        this.storedFieldLoader = fields.isEmpty() ? StoredFieldLoader.empty() : StoredFieldLoader.create(loadSource, fields);
+    }
+
+    @Override
+    public Source getSource(LeafReaderContext ctx, int doc) throws IOException {
+        final Object id = ctx.id();
+        var leaf = leaves.get(id);
+        if (leaf == null) {
+            leaf = new Leaf(sourceLoader.leaf(ctx.reader(), null), storedFieldLoader.getLoader(ctx, null));
+            var existing = leaves.put(id, leaf);
+            assert existing == null : "unexpected source provider [" + existing + "]";
+        }
+        return leaf.getSource(ctx, doc);
+    }
+
+    private static class Leaf implements SourceProvider {
+        private final SourceLoader.Leaf sourceLoader;
+        private final LeafStoredFieldLoader storedFieldLoader;
+        int doc = -1;
+        Source source = null;
+
+        private Leaf(SourceLoader.Leaf sourceLoader, LeafStoredFieldLoader storedFieldLoader) {
+            this.sourceLoader = sourceLoader;
+            this.storedFieldLoader = storedFieldLoader;
+        }
+
+        @Override
+        public Source getSource(LeafReaderContext ctx, int doc) throws IOException {
+            if (this.doc == doc) {
+                return source;
+            }
+            this.doc = doc;
+            storedFieldLoader.advanceTo(doc);
+            return source = sourceLoader.source(storedFieldLoader, doc);
+        }
+    }
+}

server/src/main/java/org/elasticsearch/search/lookup/SourceProvider.java

@@ -10,10 +10,8 @@
 package org.elasticsearch.search.lookup;
 
 import org.apache.lucene.index.LeafReaderContext;
-import org.elasticsearch.index.fieldvisitor.StoredFieldLoader;
-import org.elasticsearch.index.mapper.Mapping;
+import org.elasticsearch.index.mapper.MappingLookup;
 import org.elasticsearch.index.mapper.SourceFieldMetrics;
-import org.elasticsearch.index.mapper.SourceLoader;
 
 import java.io.IOException;
 
@@ -28,27 +26,14 @@ public interface SourceProvider {
     Source getSource(LeafReaderContext ctx, int doc) throws IOException;
 
     /**
-     * A SourceProvider that loads source from stored fields
+     * A SourceProvider that delegate loading source to the provided {@link MappingLookup}.
      *
      * The returned SourceProvider is thread-safe across segments, in that it may be
      * safely used by a searcher that searches different segments on different threads,
      * but it is not safe to use this to access documents from the same segment across
      * multiple threads.
      */
-    static SourceProvider fromStoredFields() {
-        StoredFieldLoader storedFieldLoader = StoredFieldLoader.sequentialSource();
-        return new StoredFieldSourceProvider(storedFieldLoader);
-    }
-
-    /**
-     * A SourceProvider that loads source from synthetic source
-     *
-     * The returned SourceProvider is thread-safe across segments, in that it may be
-     * safely used by a searcher that searches different segments on different threads,
-     * but it is not safe to use this to access documents from the same segment across
-     * multiple threads.
-     */
-    static SourceProvider fromSyntheticSource(Mapping mapping, SourceFieldMetrics metrics) {
-        return new SyntheticSourceProvider(new SourceLoader.Synthetic(mapping::syntheticFieldLoader, metrics));
+    static SourceProvider fromLookup(MappingLookup lookup, SourceFieldMetrics metrics) {
+        return new ConcurrentSegmentSourceProvider(lookup.newSourceLoader(metrics));
     }
 }

server/src/main/java/org/elasticsearch/search/lookup/StoredFieldSourceProvider.java

@@ -346,7 +346,7 @@ public void testParseDocumentSubFieldAccess() throws IOException {
                 (mft, lookupSupplier, fdo) -> mft.fielddataBuilder(
                     new FieldDataContext("test", null, lookupSupplier, mapperService.mappingLookup()::sourcePaths, fdo)
                 ).build(null, null),
-                SourceProvider.fromStoredFields()
+                SourceProvider.fromLookup(mapperService.mappingLookup(), mapperService.getMapperMetrics().sourceFieldMetrics())
             );
 
             LeafSearchLookup leafSearchLookup = searchLookup.getLeafSearchLookup(reader.leaves().get(0));

server/src/main/java/org/elasticsearch/search/lookup/SyntheticSourceProvider.java

@@ -106,7 +106,11 @@ public final void testFromSourceDoesNotEnforceValuesLimit() throws IOException {
                 DateFieldScript.LeafFactory leafFactory = fromSource().newFactory(
                     "field",
                     Collections.emptyMap(),
-                    new SearchLookup(field -> null, (ft, lookup, fdt) -> null, SourceProvider.fromStoredFields()),
+                    new SearchLookup(
+                        field -> null,
+                        (ft, lookup, fdt) -> null,
+                        SourceProvider.fromLookup(MappingLookup.EMPTY, SourceFieldMetrics.NOOP)
+                    ),
                     DateFormatter.forPattern("epoch_millis"),
                     OnScriptError.FAIL
                 );

server/src/test/java/org/elasticsearch/index/mapper/CompositeRuntimeFieldTests.java

@@ -103,7 +103,11 @@ public final void testFromSourceDoesNotEnforceValuesLimit() throws IOException {
                 DoubleFieldScript.LeafFactory leafFactory = fromSource().newFactory(
                     "field",
                     Collections.emptyMap(),
-                    new SearchLookup(field -> null, (ft, lookup, fdt) -> null, SourceProvider.fromStoredFields()),
+                    new SearchLookup(
+                        field -> null,
+                        (ft, lookup, fdt) -> null,
+                        SourceProvider.fromLookup(MappingLookup.EMPTY, SourceFieldMetrics.NOOP)
+                    ),
                     OnScriptError.FAIL
                 );
                 DoubleFieldScript doubleFieldScript = leafFactory.newInstance(reader.leaves().get(0));