Runtime fields: pass down runtime field name as source filter when source mode is synthetic

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;
@@ -48,18 +49,21 @@ public abstract class AbstractScriptFieldType<LeafFactory> extends MappedFieldTy
     protected final Script script;
     private final Function<SearchLookup, LeafFactory> factory;
     private final boolean isResultDeterministic;
+    private final boolean isParsedFromSource;
 
     protected AbstractScriptFieldType(
         String name,
         Function<SearchLookup, LeafFactory> factory,
         Script script,
         boolean isResultDeterministic,
-        Map<String, String> meta
+        Map<String, String> meta,
+        boolean isParsedFromSource
     ) {
         super(name, false, false, false, TextSearchInfo.SIMPLE_MATCH_WITHOUT_TERMS, meta);
         this.factory = factory;
         this.script = Objects.requireNonNull(script);
         this.isResultDeterministic = isResultDeterministic;
+        this.isParsedFromSource = isParsedFromSource;
     }
 
     @Override
@@ -190,7 +194,13 @@ public ValueFetcher valueFetcher(SearchExecutionContext context, String format)
      * Create a script leaf factory.
      */
     protected final LeafFactory leafFactory(SearchLookup searchLookup) {
-        return factory.apply(searchLookup);
+        if (isParsedFromSource) {
+            String include = name();
+            var copy = searchLookup.optimizedSourceProvider(new SourceFilter(new String[] { include }, new String[0]));
+            return factory.apply(copy);
+        } else {
+            return factory.apply(searchLookup);
+        }
     }
 
     /**

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

@@ -83,7 +83,8 @@ public static RuntimeField sourceOnly(String name) {
             searchLookup -> scriptFactory.newFactory(name, script.getParams(), searchLookup, onScriptError),
             script,
             scriptFactory.isResultDeterministic(),
-            meta
+            meta,
+            scriptFactory.isParsedFromSource()
         );
     }
 

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

@@ -148,7 +148,8 @@ public static RuntimeField sourceOnly(String name, DateFormatter dateTimeFormatt
             searchLookup -> scriptFactory.newFactory(name, script.getParams(), searchLookup, dateTimeFormatter, onScriptError),
             script,
             scriptFactory.isResultDeterministic(),
-            meta
+            meta,
+            scriptFactory.isParsedFromSource()
         );
         this.dateTimeFormatter = dateTimeFormatter;
         this.dateMathParser = dateTimeFormatter.toDateMathParser();

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

@@ -83,7 +83,8 @@ public static RuntimeField sourceOnly(String name) {
             searchLookup -> scriptFactory.newFactory(name, script.getParams(), searchLookup, onScriptError),
             script,
             scriptFactory.isResultDeterministic(),
-            meta
+            meta,
+            scriptFactory.isParsedFromSource()
         );
     }
 

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

@@ -83,7 +83,8 @@ protected GeoPointFieldScript.Factory getCompositeLeafFactory(
             searchLookup -> scriptFactory.newFactory(name, script.getParams(), searchLookup, onScriptError),
             script,
             scriptFactory.isResultDeterministic(),
-            meta
+            meta,
+            scriptFactory.isParsedFromSource()
         );
     }
 

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

@@ -81,7 +81,8 @@ protected IpFieldScript.Factory getCompositeLeafFactory(
             searchLookup -> scriptFactory.newFactory(name, script.getParams(), searchLookup, onScriptError),
             script,
             scriptFactory.isResultDeterministic(),
-            meta
+            meta,
+            scriptFactory.isParsedFromSource()
         );
     }
 

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

@@ -92,7 +92,8 @@ public KeywordScriptFieldType(
             searchLookup -> scriptFactory.newFactory(name, script.getParams(), searchLookup, onScriptError),
             script,
             scriptFactory.isResultDeterministic(),
-            meta
+            meta,
+            scriptFactory.isParsedFromSource()
         );
     }
 

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

@@ -83,7 +83,8 @@ public LongScriptFieldType(
             searchLookup -> scriptFactory.newFactory(name, script.getParams(), searchLookup, onScriptError),
             script,
             scriptFactory.isResultDeterministic(),
-            meta
+            meta,
+            scriptFactory.isParsedFromSource()
         );
     }
 

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());
     }
 
     /**

server/src/main/java/org/elasticsearch/script/BooleanFieldScript.java

@@ -37,6 +37,11 @@ public void execute() {
         public boolean isResultDeterministic() {
             return true;
         }
+
+        @Override
+        public boolean isParsedFromSource() {
+            return true;
+        }
     };
 
     public static Factory leafAdapter(Function<SearchLookup, CompositeFieldScript.LeafFactory> parentFactory) {

server/src/main/java/org/elasticsearch/script/DateFieldScript.java

@@ -42,6 +42,11 @@ public void execute() {
         public boolean isResultDeterministic() {
             return true;
         }
+
+        @Override
+        public boolean isParsedFromSource() {
+            return true;
+        }
     };
 
     public static Factory leafAdapter(Function<SearchLookup, CompositeFieldScript.LeafFactory> parentFactory) {

server/src/main/java/org/elasticsearch/script/DoubleFieldScript.java

@@ -36,6 +36,11 @@ public void execute() {
         public boolean isResultDeterministic() {
             return true;
         }
+
+        @Override
+        public boolean isParsedFromSource() {
+            return true;
+        }
     };
 
     public static Factory leafAdapter(Function<SearchLookup, CompositeFieldScript.LeafFactory> parentFactory) {

server/src/main/java/org/elasticsearch/script/GeoPointFieldScript.java

@@ -46,6 +46,11 @@ public void execute() {
         public boolean isResultDeterministic() {
             return true;
         }
+
+        @Override
+        public boolean isParsedFromSource() {
+            return true;
+        }
     };
 
     public static Factory leafAdapter(Function<SearchLookup, CompositeFieldScript.LeafFactory> parentFactory) {

server/src/main/java/org/elasticsearch/script/GeometryFieldScript.java

@@ -44,6 +44,11 @@ public void execute() {
         public boolean isResultDeterministic() {
             return true;
         }
+
+        @Override
+        public boolean isParsedFromSource() {
+            return true;
+        }
     };
 
     public static Factory leafAdapter(Function<SearchLookup, CompositeFieldScript.LeafFactory> parentFactory) {

server/src/main/java/org/elasticsearch/script/IpFieldScript.java

@@ -57,6 +57,11 @@ public void execute() {
         public boolean isResultDeterministic() {
             return true;
         }
+
+        @Override
+        public boolean isParsedFromSource() {
+            return true;
+        }
     };
 
     public static Factory leafAdapter(Function<SearchLookup, CompositeFieldScript.LeafFactory> parentFactory) {

server/src/main/java/org/elasticsearch/script/LongFieldScript.java

@@ -35,6 +35,11 @@ public void execute() {
         public boolean isResultDeterministic() {
             return true;
         }
+
+        @Override
+        public boolean isParsedFromSource() {
+            return true;
+        }
     };
 
     public static Factory leafAdapter(Function<SearchLookup, CompositeFieldScript.LeafFactory> parentFactory) {

server/src/main/java/org/elasticsearch/script/ScriptFactory.java

@@ -17,4 +17,9 @@ public interface ScriptFactory {
     default boolean isResultDeterministic() {
         return false;
     }
+
+    /** Returns {@code true} if the script only parses a field from source */
+    default boolean isParsedFromSource() {
+        return false;
+    }
 }

server/src/main/java/org/elasticsearch/script/StringFieldScript.java

@@ -43,6 +43,11 @@ public void execute() {
         public boolean isResultDeterministic() {
             return true;
         }
+
+        @Override
+        public boolean isParsedFromSource() {
+            return true;
+        }
     };
 
     public static Factory leafAdapter(Function<SearchLookup, CompositeFieldScript.LeafFactory> parentFactory) {

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,27 @@
  * 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);
         // we force a sequential reader here since it is used during query execution where documents are scanned sequentially
+        this.isStoredSource = lookup.isSourceSynthetic() == false;
+        this.storedFieldLoader = StoredFieldLoader.create(isStoredSource, sourceLoader.requiredStoredFields(), true);
+    }
+
+    private ConcurrentSegmentSourceProvider(ConcurrentSegmentSourceProvider source, SourceFilter filter) {
+        assert source.isStoredSource == false;
+        this.sourceLoaderProvider = source.sourceLoaderProvider;
+        this.isStoredSource = source.isStoredSource;
+        this.sourceLoader = source.sourceLoaderProvider.apply(filter);
+        // Also re-initialize stored field loader:
         this.storedFieldLoader = StoredFieldLoader.create(isStoredSource, sourceLoader.requiredStoredFields(), true);
-        this.isStoredSource = isStoredSource;
     }
 
     @Override
@@ -58,6 +72,15 @@ public Source getSource(LeafReaderContext ctx, int doc) throws IOException {
         return leaf.getSource(ctx, doc);
     }
 
+    @Override
+    public SourceProvider optimizedSourceProvider(SourceFilter sourceFilter) {
+        if (isStoredSource) {
+            return this;
+        } else {
+            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;
     }
 
+    private 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 optimizedSourceProvider(SourceFilter sourceFilter) {
+        SourceProvider copy = sourceProvider.optimizedSourceProvider(sourceFilter);
+        return new SearchLookup(this, copy);
+    }
 }

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

@@ -34,6 +34,19 @@ 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);
+    }
+
+    /**
+     * Optionally returns a new {@link SourceProvider} that is more optimized to load source with the provided source filter in mind.
+     * <p>
+     * Currently this is only the case if source mode is synthetic, and only a subset of fields is requested,
+     * then only loading source for requested fields is much more efficient.
+     *
+     * @param sourceFilter The part of the source caller is actually interested in.
+     * @return a new instance if source can be loaded in a more optimal way, otherwise returns this instance.
+     */
+    default SourceProvider optimizedSourceProvider(SourceFilter sourceFilter) {
+        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);
                 }
             };