diff --git a/docs/changelog/133897.yaml b/docs/changelog/133897.yaml new file mode 100644 index 0000000000000..b402284a52cc7 --- /dev/null +++ b/docs/changelog/133897.yaml @@ -0,0 +1,6 @@ +pr: 133897 +summary: "Runtime fields: pass down runtime field name as source filter when source\ + \ mode is synthetic" +area: "Mapping" +type: enhancement +issues: [] diff --git a/server/src/main/java/org/elasticsearch/index/mapper/AbstractScriptFieldType.java b/server/src/main/java/org/elasticsearch/index/mapper/AbstractScriptFieldType.java index 4de8a72f25aa9..1140a45fe0308 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/AbstractScriptFieldType.java +++ b/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 extends MappedFieldTy protected final Script script; private final Function factory; private final boolean isResultDeterministic; + private final boolean isParsedFromSource; protected AbstractScriptFieldType( String name, Function factory, Script script, boolean isResultDeterministic, - Map meta + Map 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); + } } /** diff --git a/server/src/main/java/org/elasticsearch/index/mapper/BooleanScriptFieldType.java b/server/src/main/java/org/elasticsearch/index/mapper/BooleanScriptFieldType.java index dcddda983866f..b22ef6333ede1 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/BooleanScriptFieldType.java +++ b/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() ); } diff --git a/server/src/main/java/org/elasticsearch/index/mapper/DateScriptFieldType.java b/server/src/main/java/org/elasticsearch/index/mapper/DateScriptFieldType.java index 341944c3d687a..d67d8792dcf9a 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/DateScriptFieldType.java +++ b/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(); diff --git a/server/src/main/java/org/elasticsearch/index/mapper/DoubleScriptFieldType.java b/server/src/main/java/org/elasticsearch/index/mapper/DoubleScriptFieldType.java index 1eac55bae51a8..7c7cd2edca6d4 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/DoubleScriptFieldType.java +++ b/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() ); } diff --git a/server/src/main/java/org/elasticsearch/index/mapper/GeoPointScriptFieldType.java b/server/src/main/java/org/elasticsearch/index/mapper/GeoPointScriptFieldType.java index 1aeb32f4d2b73..cf6aa6b521308 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/GeoPointScriptFieldType.java +++ b/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() ); } diff --git a/server/src/main/java/org/elasticsearch/index/mapper/IpScriptFieldType.java b/server/src/main/java/org/elasticsearch/index/mapper/IpScriptFieldType.java index eb827f257515b..9923235f09025 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/IpScriptFieldType.java +++ b/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() ); } diff --git a/server/src/main/java/org/elasticsearch/index/mapper/KeywordScriptFieldType.java b/server/src/main/java/org/elasticsearch/index/mapper/KeywordScriptFieldType.java index e1778b19765f5..d0f5b2b61edae 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/KeywordScriptFieldType.java +++ b/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() ); } diff --git a/server/src/main/java/org/elasticsearch/index/mapper/LongScriptFieldType.java b/server/src/main/java/org/elasticsearch/index/mapper/LongScriptFieldType.java index 513daf198f167..9cbf60dd87240 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/LongScriptFieldType.java +++ b/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() ); } diff --git a/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java b/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java index 56e136801e128..6e0bf8414d84d 100644 --- a/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java +++ b/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()); } /** diff --git a/server/src/main/java/org/elasticsearch/script/BooleanFieldScript.java b/server/src/main/java/org/elasticsearch/script/BooleanFieldScript.java index f1730014921e4..841e3c336c594 100644 --- a/server/src/main/java/org/elasticsearch/script/BooleanFieldScript.java +++ b/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 parentFactory) { diff --git a/server/src/main/java/org/elasticsearch/script/DateFieldScript.java b/server/src/main/java/org/elasticsearch/script/DateFieldScript.java index 9aadfdbb79bf9..79d3f6e0010d7 100644 --- a/server/src/main/java/org/elasticsearch/script/DateFieldScript.java +++ b/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 parentFactory) { diff --git a/server/src/main/java/org/elasticsearch/script/DoubleFieldScript.java b/server/src/main/java/org/elasticsearch/script/DoubleFieldScript.java index 9477c1998725e..b262f786da563 100644 --- a/server/src/main/java/org/elasticsearch/script/DoubleFieldScript.java +++ b/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 parentFactory) { diff --git a/server/src/main/java/org/elasticsearch/script/GeoPointFieldScript.java b/server/src/main/java/org/elasticsearch/script/GeoPointFieldScript.java index bbf703e8f1930..cbea042cc16ae 100644 --- a/server/src/main/java/org/elasticsearch/script/GeoPointFieldScript.java +++ b/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 parentFactory) { diff --git a/server/src/main/java/org/elasticsearch/script/GeometryFieldScript.java b/server/src/main/java/org/elasticsearch/script/GeometryFieldScript.java index ffd0d9690d20c..3c5312cf2057f 100644 --- a/server/src/main/java/org/elasticsearch/script/GeometryFieldScript.java +++ b/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 parentFactory) { diff --git a/server/src/main/java/org/elasticsearch/script/IpFieldScript.java b/server/src/main/java/org/elasticsearch/script/IpFieldScript.java index 51120bfefaad0..b7a5526d8b1c5 100644 --- a/server/src/main/java/org/elasticsearch/script/IpFieldScript.java +++ b/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 parentFactory) { diff --git a/server/src/main/java/org/elasticsearch/script/LongFieldScript.java b/server/src/main/java/org/elasticsearch/script/LongFieldScript.java index c600cf40b8417..9a577ce142c13 100644 --- a/server/src/main/java/org/elasticsearch/script/LongFieldScript.java +++ b/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 parentFactory) { diff --git a/server/src/main/java/org/elasticsearch/script/ScriptFactory.java b/server/src/main/java/org/elasticsearch/script/ScriptFactory.java index 6ddc05267cefb..2adc373c43ca2 100644 --- a/server/src/main/java/org/elasticsearch/script/ScriptFactory.java +++ b/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; + } } diff --git a/server/src/main/java/org/elasticsearch/script/StringFieldScript.java b/server/src/main/java/org/elasticsearch/script/StringFieldScript.java index 4d43baf859332..f3a61df5bf40e 100644 --- a/server/src/main/java/org/elasticsearch/script/StringFieldScript.java +++ b/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 parentFactory) { diff --git a/server/src/main/java/org/elasticsearch/search/lookup/ConcurrentSegmentSourceProvider.java b/server/src/main/java/org/elasticsearch/search/lookup/ConcurrentSegmentSourceProvider.java index 13c406d7cceeb..1b7ddbeea8d05 100644 --- a/server/src/main/java/org/elasticsearch/search/lookup/ConcurrentSegmentSourceProvider.java +++ b/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 sourceLoaderProvider; private final SourceLoader sourceLoader; private final StoredFieldLoader storedFieldLoader; private final Map 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; diff --git a/server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java b/server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java index d899a390d8be7..4e932ea47ca7e 100644 --- a/server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java +++ b/server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java @@ -102,6 +102,14 @@ private SearchLookup(SearchLookup searchLookup, Set 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); + } } diff --git a/server/src/main/java/org/elasticsearch/search/lookup/SourceProvider.java b/server/src/main/java/org/elasticsearch/search/lookup/SourceProvider.java index 954930a9437c9..b655607fc8b38 100644 --- a/server/src/main/java/org/elasticsearch/search/lookup/SourceProvider.java +++ b/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. + *

+ * 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 the 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; } } diff --git a/server/src/test/java/org/elasticsearch/index/mapper/RuntimeFieldSourceProviderOptimizationTests.java b/server/src/test/java/org/elasticsearch/index/mapper/RuntimeFieldSourceProviderOptimizationTests.java new file mode 100644 index 0000000000000..03234f12e421c --- /dev/null +++ b/server/src/test/java/org/elasticsearch/index/mapper/RuntimeFieldSourceProviderOptimizationTests.java @@ -0,0 +1,177 @@ +/* + * 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.index.mapper; + +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.store.Directory; +import org.elasticsearch.common.bytes.BytesArray; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.index.IndexSettings; +import org.elasticsearch.script.LongFieldScript; +import org.elasticsearch.search.lookup.SearchLookup; +import org.elasticsearch.test.ESSingleNodeTestCase; +import org.elasticsearch.xcontent.XContentType; + +import java.io.IOException; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder; +import static org.hamcrest.Matchers.equalTo; + +/** + * Tests that source provider optimization that filters _source based on the same of source only runtime fields kick in. + * This is important for synthetic source, otherwise many doc value and stored fields get loaded in the process. + */ +public class RuntimeFieldSourceProviderOptimizationTests extends ESSingleNodeTestCase { + + public void testWithSourceProviderOptimization() throws IOException { + var mapping = jsonBuilder().startObject().startObject("runtime").startObject("field"); + mapping.field("type", "long"); + mapping.endObject().endObject().endObject(); + var indexService = createIndex("test-index", Settings.builder().put("index.mapping.source.mode", "synthetic").build(), mapping); + + int numDocs = 256; + try (Directory directory = newDirectory(); IndexWriter iw = new IndexWriter(directory, new IndexWriterConfig())) { + for (int i = 0; i < numDocs; i++) { + BytesArray source = new BytesArray(String.format(Locale.ROOT, "{\"field\":%d,\"another_field\":123}", i)); + var doc = indexService.mapperService() + .documentMapper() + .parse(new SourceToParse(Integer.toString(i), source, XContentType.JSON)) + .rootDoc(); + iw.addDocument(doc); + } + iw.commit(); + iw.forceMerge(1); + + try (var indexReader = DirectoryReader.open(iw)) { + var searcher = new IndexSearcher(indexReader); + LeafReaderContext leafReaderContext = indexReader.leaves().getFirst(); + var context = indexService.newSearchExecutionContext(0, 0, searcher, () -> 1L, null, Map.of()); + var fieldType = (AbstractScriptFieldType) indexService.mapperService().fieldType("field"); + + // The other_field should have been filtered out, otherwise the mechanism that pushes field name as source filter to + // SourceProvider isn't kicking in. Essentially checking that optimization in + // ConcurrentSegmentSourceProvider.optimizedSourceProvider(...) kicks in: + var leafFactory = (LongFieldScript.LeafFactory) fieldType.leafFactory(context); + var fieldScript = leafFactory.newInstance(leafReaderContext); + for (int i = 0; i < 256; i++) { + fieldScript.runForDoc(i); + var source = fieldScript.source().get().source(); + assertThat(source, equalTo(Map.of("field", i))); + } + + // Test that runtime based term query works as expected with the optimization: + var termQuery = fieldType.termQuery(32, context); + assertThat(searcher.count(termQuery), equalTo(1)); + + // Test that runtime based block loader works as expected with the optimization: + var blockLoader = fieldType.blockLoader(blContext(context.lookup())); + var columnReader = blockLoader.columnAtATimeReader(leafReaderContext); + var block = (TestBlock) columnReader.read(TestBlock.factory(), TestBlock.docs(leafReaderContext), 0, false); + for (int i = 0; i < block.size(); i++) { + assertThat(block.get(i), equalTo((long) i)); + } + } + } + } + + public void testWithoutSourceProviderOptimization() throws IOException { + var mapping = jsonBuilder().startObject().startObject("runtime").startObject("field"); + mapping.field("type", "long"); + mapping.endObject().endObject().endObject(); + var indexService = createIndex("test-index", Settings.EMPTY, mapping); + + int numDocs = 256; + try (Directory directory = newDirectory(); IndexWriter iw = new IndexWriter(directory, new IndexWriterConfig())) { + for (int i = 0; i < numDocs; i++) { + BytesArray source = new BytesArray(String.format(Locale.ROOT, "{\"field\":%d,\"another_field\":123}", i)); + var doc = indexService.mapperService() + .documentMapper() + .parse(new SourceToParse(Integer.toString(i), source, XContentType.JSON)) + .rootDoc(); + iw.addDocument(doc); + } + iw.commit(); + iw.forceMerge(1); + + try (var indexReader = DirectoryReader.open(iw)) { + var searcher = new IndexSearcher(indexReader); + LeafReaderContext leafReaderContext = indexReader.leaves().getFirst(); + var context = indexService.newSearchExecutionContext(0, 0, searcher, () -> 1L, null, Map.of()); + var fieldType = (AbstractScriptFieldType) indexService.mapperService().fieldType("field"); + + var leafFactory = (LongFieldScript.LeafFactory) fieldType.leafFactory(context); + var fieldScript = leafFactory.newInstance(leafReaderContext); + for (int i = 0; i < 256; i++) { + fieldScript.runForDoc(i); + var source = fieldScript.source().get().source(); + assertThat(source, equalTo(Map.of("field", i, "another_field", 123))); + } + + // Test that runtime based term query works as expected with the optimization: + var termQuery = fieldType.termQuery(32, context); + assertThat(searcher.count(termQuery), equalTo(1)); + + // Test that runtime based block loader works as expected with the optimization: + var blockLoader = fieldType.blockLoader(blContext(context.lookup())); + var columnReader = blockLoader.columnAtATimeReader(leafReaderContext); + var block = (TestBlock) columnReader.read(TestBlock.factory(), TestBlock.docs(leafReaderContext), 0, false); + for (int i = 0; i < block.size(); i++) { + assertThat(block.get(i), equalTo((long) i)); + } + } + } + } + + static MappedFieldType.BlockLoaderContext blContext(SearchLookup lookup) { + return new MappedFieldType.BlockLoaderContext() { + @Override + public String indexName() { + throw new UnsupportedOperationException(); + } + + @Override + public IndexSettings indexSettings() { + throw new UnsupportedOperationException(); + } + + @Override + public MappedFieldType.FieldExtractPreference fieldExtractPreference() { + return MappedFieldType.FieldExtractPreference.NONE; + } + + @Override + public SearchLookup lookup() { + return lookup; + } + + @Override + public Set sourcePaths(String name) { + throw new UnsupportedOperationException(); + } + + @Override + public String parentField(String field) { + throw new UnsupportedOperationException(); + } + + @Override + public FieldNamesFieldMapper.FieldNamesFieldType fieldNames() { + return FieldNamesFieldMapper.FieldNamesFieldType.get(true); + } + }; + } +} diff --git a/test/framework/src/main/java/org/elasticsearch/index/mapper/AbstractScriptFieldTypeTestCase.java b/test/framework/src/main/java/org/elasticsearch/index/mapper/AbstractScriptFieldTypeTestCase.java index fa35a4ee1a4da..f679d47274b46 100644 --- a/test/framework/src/main/java/org/elasticsearch/index/mapper/AbstractScriptFieldTypeTestCase.java +++ b/test/framework/src/main/java/org/elasticsearch/index/mapper/AbstractScriptFieldTypeTestCase.java @@ -55,6 +55,9 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; public abstract class AbstractScriptFieldTypeTestCase extends MapperServiceTestCase { @@ -420,6 +423,46 @@ public final void testCacheable() throws IOException { } } + public final void testIsParsedFromSource() throws IOException { + XContentBuilder mapping = runtimeMapping(b -> { + b.startObject("field") + .field("type", typeName()) + .startObject("script") + .field("source", "dummy_source") + .field("lang", "test") + .endObject() + .endObject() + .startObject("field_source") + .field("type", typeName()) + .startObject("script") + .field("source", "deterministic_source") + .field("lang", "test") + .endObject() + .endObject(); + }); + MapperService mapperService = createMapperService(mapping); + SearchExecutionContext c = createSearchExecutionContext(mapperService); + { + // The field_source uses parseFromSource(...) in compileScript(...) method in this class. + // This triggers calling SearchLookup#optimizedSourceProvider(...) which should return more optimized source. + var fieldType = (AbstractScriptFieldType) c.getFieldType("field_source"); + SearchLookup searchLookup = mock(SearchLookup.class); + when(searchLookup.optimizedSourceProvider(any())).thenReturn(searchLookup); + var result = fieldType.leafFactory(searchLookup); + assertNotNull(result); + verify(searchLookup, times(1)).optimizedSourceProvider(any()); + } + { + // The field uses normal scripts and that should never cause SearchLookup#optimizedSourceProvider(...) to be invoked: + var fieldType = (AbstractScriptFieldType) c.getFieldType("field"); + SearchLookup searchLookup = mock(SearchLookup.class); + when(searchLookup.optimizedSourceProvider(any())).thenReturn(searchLookup); + var result = fieldType.leafFactory(searchLookup); + assertNotNull(result); + verify(searchLookup, never()).optimizedSourceProvider(any()); + } + } + protected final List blockLoaderReadValuesFromColumnAtATimeReader(DirectoryReader reader, MappedFieldType fieldType, int offset) throws IOException { BlockLoader loader = fieldType.blockLoader(blContext()); diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java index 19012e41ecdc2..dddb41e9510c2 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java +++ b/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); } }; diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ReinitializingSourceProvider.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ReinitializingSourceProvider.java index 61ac67674b252..5e60033fbabd9 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ReinitializingSourceProvider.java +++ b/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 sourceProviderFactory; + private final Function sourceProviderFactory; - ReinitializingSourceProvider(Supplier sourceProviderFactory) { + ReinitializingSourceProvider(Function sourceProviderFactory) { + this.sourceFilter = null; + this.sourceProviderFactory = sourceProviderFactory; + } + + private ReinitializingSourceProvider(SourceFilter sourceFilter, Function sourceProviderFactory) { + this.sourceFilter = sourceFilter; this.sourceProviderFactory = sourceProviderFactory; } @@ -37,13 +45,18 @@ 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 optimizedSourceProvider(SourceFilter sourceFilter) { + return new ReinitializingSourceProvider(sourceFilter, sourceProviderFactory); + } + private static final class PerThreadSourceProvider { final SourceProvider source; final Thread creatingThread; diff --git a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeScriptFieldType.java b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeScriptFieldType.java index 358f6768132b9..c11f77b8da74e 100644 --- a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeScriptFieldType.java +++ b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeScriptFieldType.java @@ -82,7 +82,8 @@ protected GeometryFieldScript.Factory getCompositeLeafFactory( searchLookup -> scriptFactory.newFactory(name, script.getParams(), searchLookup, onScriptError), script, scriptFactory.isResultDeterministic(), - meta + meta, + scriptFactory.isParsedFromSource() ); this.geoFormatterFactory = geoFormatterFactory; }