Refactor SourceProvider creation to consistently use MappingLookup

This change updates the code to always create SourceProvider instances via MappingLookup, avoiding direct exposure to the underlying source format (synthetic or stored).
It also aligns source filtering behavior between SourceProvider and SourceLoader, ensuring consistent application of filters.

This change is needed to enable source filtering to occur earlier in the fetch phase, for example, when constructing a synthetic source.

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.PluginsLoader;
@@ -90,7 +92,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, null, 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(), null, 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/FilteredSearchExecutionContext.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.IndexVersion;
@@ -39,6 +40,7 @@
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.search.lookup.LeafFieldLookupProvider;
 import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.search.lookup.SourceFilter;
 import org.elasticsearch.search.lookup.SourceProvider;
 import org.elasticsearch.xcontent.XContentParserConfiguration;
 
@@ -162,8 +164,8 @@ public boolean isSourceSynthetic() {
     }
 
     @Override
-    public SourceLoader newSourceLoader(boolean forceSyntheticSource) {
-        return in.newSourceLoader(forceSyntheticSource);
+    public SourceLoader newSourceLoader(@Nullable SourceFilter filter, boolean forceSyntheticSource) {
+        return in.newSourceLoader(filter, forceSyntheticSource);
     }
 
     @Override

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

@@ -25,6 +25,7 @@
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.lucene.search.Queries;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.IndexSortConfig;
@@ -57,6 +58,7 @@
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.search.lookup.LeafFieldLookupProvider;
 import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.search.lookup.SourceFilter;
 import org.elasticsearch.search.lookup.SourceProvider;
 import org.elasticsearch.transport.RemoteClusterAware;
 import org.elasticsearch.xcontent.XContentParserConfiguration;
@@ -439,15 +441,15 @@ public boolean isSourceSynthetic() {
     /**
      * Build something to load source {@code _source}.
      */
-    public SourceLoader newSourceLoader(boolean forceSyntheticSource) {
+    public SourceLoader newSourceLoader(@Nullable SourceFilter filter, boolean forceSyntheticSource) {
         if (forceSyntheticSource) {
             return new SourceLoader.Synthetic(
-                null,
+                filter,
                 () -> mappingLookup.getMapping().syntheticFieldLoader(null),
                 mapperMetrics.sourceFieldMetrics()
             );
         }
-        return mappingLookup.newSourceLoader(null, mapperMetrics.sourceFieldMetrics());
+        return mappingLookup.newSourceLoader(filter, mapperMetrics.sourceFieldMetrics());
     }
 
     /**
@@ -506,9 +508,7 @@ public SearchLookup lookup() {
     }
 
     public SourceProvider createSourceProvider() {
-        return isSourceSynthetic()
-            ? SourceProvider.fromSyntheticSource(mappingLookup.getMapping(), null, mapperMetrics.sourceFieldMetrics())
-            : SourceProvider.fromStoredFields();
+        return SourceProvider.fromLookup(mappingLookup, null, mapperMetrics.sourceFieldMetrics());
     }
 
     /**

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

@@ -70,6 +70,7 @@
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.internal.ShardSearchContextId;
 import org.elasticsearch.search.internal.ShardSearchRequest;
+import org.elasticsearch.search.lookup.SourceFilter;
 import org.elasticsearch.search.profile.Profilers;
 import org.elasticsearch.search.query.QuerySearchResult;
 import org.elasticsearch.search.rank.context.QueryPhaseRankShardContext;
@@ -943,8 +944,8 @@ public ReaderContext readerContext() {
     }
 
     @Override
-    public SourceLoader newSourceLoader() {
-        return searchExecutionContext.newSourceLoader(request.isForceSyntheticSource());
+    public SourceLoader newSourceLoader(SourceFilter filter) {
+        return searchExecutionContext.newSourceLoader(filter, request.isForceSyntheticSource());
     }
 
     @Override

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

@@ -45,9 +45,9 @@ public class FetchContext {
     /**
      * Create a FetchContext based on a SearchContext
      */
-    public FetchContext(SearchContext searchContext) {
+    public FetchContext(SearchContext searchContext, SourceLoader sourceLoader) {
         this.searchContext = searchContext;
-        this.sourceLoader = searchContext.newSourceLoader();
+        this.sourceLoader = sourceLoader;
         this.storedFieldsContext = buildStoredFieldsContext(searchContext);
         this.fetchSourceContext = buildFetchSourceContext(searchContext);
     }

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

@@ -111,9 +111,8 @@ public Source getSource(LeafReaderContext ctx, int doc) {
     }
 
     private SearchHits buildSearchHits(SearchContext context, int[] docIdsToLoad, Profiler profiler, RankDocShardInfo rankDocs) {
-
-        FetchContext fetchContext = new FetchContext(context);
-        SourceLoader sourceLoader = context.newSourceLoader();
+        SourceLoader sourceLoader = context.newSourceLoader(null);
+        FetchContext fetchContext = new FetchContext(context, sourceLoader);
 
         PreloadedSourceProvider sourceProvider = new PreloadedSourceProvider();
         PreloadedFieldLookupProvider fieldLookupProvider = new PreloadedFieldLookupProvider();

server/src/main/java/org/elasticsearch/search/internal/FilteredSearchContext.java

@@ -14,6 +14,7 @@
 import org.apache.lucene.search.TotalHits;
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.cache.bitset.BitsetFilterCache;
 import org.elasticsearch.index.mapper.IdLoader;
@@ -33,6 +34,7 @@
 import org.elasticsearch.search.fetch.subphase.InnerHitsContext;
 import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext;
 import org.elasticsearch.search.fetch.subphase.highlight.SearchHighlightContext;
+import org.elasticsearch.search.lookup.SourceFilter;
 import org.elasticsearch.search.profile.Profilers;
 import org.elasticsearch.search.query.QuerySearchResult;
 import org.elasticsearch.search.rank.context.QueryPhaseRankShardContext;
@@ -453,8 +455,8 @@ public ReaderContext readerContext() {
     }
 
     @Override
-    public SourceLoader newSourceLoader() {
-        return in.newSourceLoader();
+    public SourceLoader newSourceLoader(@Nullable SourceFilter filter) {
+        return in.newSourceLoader(filter);
     }
 
     @Override

server/src/main/java/org/elasticsearch/search/internal/SearchContext.java

@@ -41,6 +41,7 @@
 import org.elasticsearch.search.fetch.subphase.InnerHitsContext;
 import org.elasticsearch.search.fetch.subphase.ScriptFieldsContext;
 import org.elasticsearch.search.fetch.subphase.highlight.SearchHighlightContext;
+import org.elasticsearch.search.lookup.SourceFilter;
 import org.elasticsearch.search.profile.Profilers;
 import org.elasticsearch.search.query.QueryPhase;
 import org.elasticsearch.search.query.QuerySearchResult;
@@ -441,7 +442,7 @@ public String toString() {
     /**
      * Build something to load source {@code _source}.
      */
-    public abstract SourceLoader newSourceLoader();
+    public abstract SourceLoader newSourceLoader(@Nullable SourceFilter sourceFilter);
 
     public abstract IdLoader newIdLoader();
 }

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

@@ -0,0 +1,71 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", 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", the "GNU Affero General Public
+ * License v3.0 only", 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.SourceLoader;
+
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ * A {@link SourceProvider} that loads _source from a concurrent search.
+ *
+ * NOTE: 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, boolean loadSource) {
+        this.sourceLoader = loader;
+        this.storedFieldLoader = StoredFieldLoader.create(loadSource, sourceLoader.requiredStoredFields());
+    }
+
+    @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);
+        }
+    }
+}