Pass down runtime field name as source filter when source mode is synthetic

This avoids synthesizing fields that are not requested

Author: martijnvg

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

@@ -27,6 +27,7 @@
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.search.fetch.StoredFieldsSpec;
 import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.search.lookup.SourceFilter;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.time.ZoneId;
@@ -190,7 +191,9 @@ public ValueFetcher valueFetcher(SearchExecutionContext context, String format)
      * Create a script leaf factory.
      */
     protected final LeafFactory leafFactory(SearchLookup searchLookup) {
-        return factory.apply(searchLookup);
+        String include = name();
+        var copy = searchLookup.maybeCopyWithSourceFilter(new SourceFilter(new String[] { include }, new String[0]));
+        return factory.apply(copy);
     }
 
     /**

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

@@ -503,14 +503,14 @@ public boolean containsBrokenAnalysis(String field) {
      */
     public SearchLookup lookup() {
         if (this.lookup == null) {
-            var sourceProvider = createSourceProvider();
+            var sourceProvider = createSourceProvider(null);
             setLookupProviders(sourceProvider, LeafFieldLookupProvider.fromStoredFields());
         }
         return this.lookup;
     }
 
-    public SourceProvider createSourceProvider() {
-        return SourceProvider.fromLookup(mappingLookup, null, mapperMetrics.sourceFieldMetrics());
+    public SourceProvider createSourceProvider(SourceFilter sourceFilter) {
+        return SourceProvider.fromLookup(mappingLookup, sourceFilter, mapperMetrics.sourceFieldMetrics());
     }
 
     /**
@@ -543,6 +543,12 @@ public void setLookupProviders(
         );
     }
 
+    public SearchLookup copyWithSourceFilter(SourceFilter sourceFilter) {
+        var searchLookup = lookup();
+        var sourceProvider = SourceProvider.fromLookup(mappingLookup, sourceFilter, mapperMetrics.sourceFieldMetrics());
+        return new SearchLookup(searchLookup, sourceProvider);
+    }
+
     public NestedScope nestedScope() {
         return nestedScope;
     }

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

@@ -13,10 +13,13 @@
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.index.fieldvisitor.LeafStoredFieldLoader;
 import org.elasticsearch.index.fieldvisitor.StoredFieldLoader;
+import org.elasticsearch.index.mapper.MappingLookup;
+import org.elasticsearch.index.mapper.SourceFieldMetrics;
 import org.elasticsearch.index.mapper.SourceLoader;
 
 import java.io.IOException;
 import java.util.Map;
+import java.util.function.Function;
 
 /**
  * A {@link SourceProvider} that loads _source from a concurrent search.
@@ -26,16 +29,31 @@
  * within-segment concurrency this will have to work entirely differently.
  * **/
 class ConcurrentSegmentSourceProvider implements SourceProvider {
+    private final Function<SourceFilter, SourceLoader> sourceLoaderProvider;
     private final SourceLoader sourceLoader;
     private final StoredFieldLoader storedFieldLoader;
     private final Map<Object, Leaf> leaves = ConcurrentCollections.newConcurrentMap();
     private final boolean isStoredSource;
 
-    ConcurrentSegmentSourceProvider(SourceLoader loader, boolean isStoredSource) {
-        this.sourceLoader = loader;
+    ConcurrentSegmentSourceProvider(MappingLookup lookup, SourceFilter filter, SourceFieldMetrics metrics) {
+        this.sourceLoaderProvider = sourceFilter -> lookup.newSourceLoader(sourceFilter, metrics);
+        this.sourceLoader = sourceLoaderProvider.apply(filter);
+        this.isStoredSource = lookup.isSourceSynthetic() == false;
         // we force a sequential reader here since it is used during query execution where documents are scanned sequentially
         this.storedFieldLoader = StoredFieldLoader.create(isStoredSource, sourceLoader.requiredStoredFields(), true);
-        this.isStoredSource = isStoredSource;
+    }
+
+    private ConcurrentSegmentSourceProvider(ConcurrentSegmentSourceProvider source, SourceFilter filter) {
+        this.sourceLoaderProvider = source.sourceLoaderProvider;
+        this.isStoredSource = source.isStoredSource;
+        if (isStoredSource) {
+            this.sourceLoader = source.sourceLoader;
+            this.storedFieldLoader = source.storedFieldLoader;
+        } else {
+            this.sourceLoader = source.sourceLoaderProvider.apply(filter);
+            // Also re-initialize stored field loader:
+            this.storedFieldLoader = StoredFieldLoader.create(isStoredSource, sourceLoader.requiredStoredFields(), true);
+        }
     }
 
     @Override
@@ -58,6 +76,12 @@ public Source getSource(LeafReaderContext ctx, int doc) throws IOException {
         return leaf.getSource(ctx, doc);
     }
 
+    @Override
+    public SourceProvider maybeCopyWithSourceFilter(SourceFilter sourceFilter) {
+        assert leaves.isEmpty() : "source provider must be unused when applying filter";
+        return new ConcurrentSegmentSourceProvider(this, sourceFilter);
+    }
+
     private static class Leaf implements SourceProvider {
         private final SourceLoader.Leaf sourceLoader;
         private final LeafStoredFieldLoader storedFieldLoader;

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

@@ -102,6 +102,14 @@ private SearchLookup(SearchLookup searchLookup, Set<String> fieldChain) {
         this.fieldLookupProvider = searchLookup.fieldLookupProvider;
     }
 
+    public SearchLookup(SearchLookup searchLookup, SourceProvider sourceProvider) {
+        this.fieldChain = searchLookup.fieldChain;
+        this.sourceProvider = sourceProvider;
+        this.fieldTypeLookup = searchLookup.fieldTypeLookup;
+        this.fieldDataLookup = searchLookup.fieldDataLookup;
+        this.fieldLookupProvider = searchLookup.fieldLookupProvider;
+    }
+
     /**
      * Creates a copy of the current {@link SearchLookup} that looks fields up in the same way, but also tracks field references
      * in order to detect cycles and prevent resolving fields that depend on more than {@link #MAX_FIELD_CHAIN_DEPTH} other fields.
@@ -145,4 +153,8 @@ public Source getSource(LeafReaderContext ctx, int doc) throws IOException {
         return sourceProvider.getSource(ctx, doc);
     }
 
+    public SearchLookup maybeCopyWithSourceFilter(SourceFilter sourceFilter) {
+        SourceProvider copy = sourceProvider.maybeCopyWithSourceFilter(sourceFilter);
+        return new SearchLookup(this, copy);
+    }
 }

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

@@ -34,6 +34,11 @@ public interface SourceProvider {
      * multiple threads.
      */
     static SourceProvider fromLookup(MappingLookup lookup, SourceFilter filter, SourceFieldMetrics metrics) {
-        return new ConcurrentSegmentSourceProvider(lookup.newSourceLoader(filter, metrics), lookup.isSourceSynthetic() == false);
+        return new ConcurrentSegmentSourceProvider(lookup, filter, metrics);
+    }
+
+    default SourceProvider maybeCopyWithSourceFilter(SourceFilter sourceFilter) {
+        assert false : "should not be invoked";
+        return this;
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java

@@ -35,6 +35,7 @@
 import org.elasticsearch.logging.Logger;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.search.lookup.SourceFilter;
 import org.elasticsearch.search.lookup.SourceProvider;
 import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
@@ -596,7 +597,7 @@ void runCompute(CancellableTask task, ComputeContext context, PhysicalPlan plan,
             var searchExecutionContext = new SearchExecutionContext(searchContext.getSearchExecutionContext()) {
 
                 @Override
-                public SourceProvider createSourceProvider() {
+                public SourceProvider createSourceProvider(SourceFilter sourceFilter) {
                     return new ReinitializingSourceProvider(super::createSourceProvider);
                 }
             };

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ReinitializingSourceProvider.java

@@ -9,10 +9,11 @@
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.elasticsearch.search.lookup.Source;
+import org.elasticsearch.search.lookup.SourceFilter;
 import org.elasticsearch.search.lookup.SourceProvider;
 
 import java.io.IOException;
-import java.util.function.Supplier;
+import java.util.function.Function;
 
 /**
  * This class exists as a workaround for using SourceProvider in the compute engine.
@@ -25,10 +26,17 @@
  */
 final class ReinitializingSourceProvider implements SourceProvider {
 
+    private final SourceFilter sourceFilter;
     private PerThreadSourceProvider perThreadProvider;
-    private final Supplier<SourceProvider> sourceProviderFactory;
+    private final Function<SourceFilter, SourceProvider> sourceProviderFactory;
 
-    ReinitializingSourceProvider(Supplier<SourceProvider> sourceProviderFactory) {
+    ReinitializingSourceProvider(Function<SourceFilter, SourceProvider> sourceProviderFactory) {
+        this.sourceFilter = null;
+        this.sourceProviderFactory = sourceProviderFactory;
+    }
+
+    private ReinitializingSourceProvider(SourceFilter sourceFilter, Function<SourceFilter, SourceProvider> sourceProviderFactory) {
+        this.sourceFilter = sourceFilter;
         this.sourceProviderFactory = sourceProviderFactory;
     }
 
@@ -37,13 +45,19 @@ public Source getSource(LeafReaderContext ctx, int doc) throws IOException {
         var currentThread = Thread.currentThread();
         PerThreadSourceProvider provider = perThreadProvider;
         if (provider == null || provider.creatingThread != currentThread || doc < provider.lastSeenDocId) {
-            provider = new PerThreadSourceProvider(sourceProviderFactory.get(), currentThread);
+            provider = new PerThreadSourceProvider(sourceProviderFactory.apply(sourceFilter), currentThread);
             this.perThreadProvider = provider;
         }
         provider.lastSeenDocId = doc;
         return provider.source.getSource(ctx, doc);
     }
 
+    @Override
+    public SourceProvider maybeCopyWithSourceFilter(SourceFilter sourceFilter) {
+        assert perThreadProvider == null : "source provider must be unused when applying filter";
+        return new ReinitializingSourceProvider(sourceFilter, sourceProviderFactory);
+    }
+
     private static final class PerThreadSourceProvider {
         final SourceProvider source;
         final Thread creatingThread;