Improve block loader for source only runtime fields of type long. (#134117)

By using the FallbackSyntheticSourceBlockLoader instead of generic LongScriptBlockLoader.

Change RuntimeFieldSourceProviderOptimizationTests to no longer test block loading, but field data loading instead. In the latter case, the optimization that this test is testing will remain.

Also deal with the situation that a field is both a normal field and runtime field. In that case fall back to script based block loading.

Author: martijnvg

docs/changelog/134117.yaml

@@ -0,0 +1,5 @@
+pr: 134117
+summary: Improve block loader for source only runtime fields of type long
+area: Mapping
+type: enhancement
+issues: []

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

@@ -49,7 +49,7 @@ 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 final boolean isParsedFromSource;
 
     protected AbstractScriptFieldType(
         String name,

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

@@ -26,6 +26,7 @@
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.indices.breaker.NoneCircuitBreakerService;
 import org.elasticsearch.plugins.internal.XContentMeteringParserDecorator;
+import org.elasticsearch.search.lookup.LeafFieldLookupProvider;
 import org.elasticsearch.search.lookup.SearchLookup;
 import org.elasticsearch.search.lookup.Source;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -171,6 +172,7 @@ private static void executeIndexTimeScripts(DocumentParserContext context) {
         }
         SearchLookup searchLookup = new SearchLookup(
             context.mappingLookup().indexTimeLookup()::get,
+            fieldName -> context.mappingLookup().getMapper(fieldName) == null,
             (ft, lookup, fto) -> ft.fielddataBuilder(
                 new FieldDataContext(
                     context.indexSettings().getIndex().getName(),
@@ -180,7 +182,8 @@ private static void executeIndexTimeScripts(DocumentParserContext context) {
                     fto
                 )
             ).build(new IndexFieldDataCache.None(), new NoneCircuitBreakerService()),
-            (ctx, doc) -> Source.fromBytes(context.sourceToParse().source())
+            (ctx, doc) -> Source.fromBytes(context.sourceToParse().source()),
+            LeafFieldLookupProvider.fromStoredFields()
         );
         // field scripts can be called both by the loop at the end of this method and via
         // the document reader, so to ensure that we don't run them multiple times we

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

@@ -30,6 +30,7 @@
 
 import java.time.ZoneId;
 import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Function;
@@ -109,7 +110,35 @@ public DocValueFormat docValueFormat(String format, ZoneId timeZone) {
 
     @Override
     public BlockLoader blockLoader(BlockLoaderContext blContext) {
-        return new LongScriptBlockDocValuesReader.LongScriptBlockLoader(leafFactory(blContext.lookup()));
+        var indexSettings = blContext.indexSettings();
+        if (isParsedFromSource && indexSettings.getIndexMappingSourceMode() == SourceFieldMapper.Mode.SYNTHETIC
+        // A runtime and normal field can share the same name.
+        // In that case there is no ignored source entry, and so we need to fail back to LongScriptBlockLoader.
+        // We could optimize this, but at this stage feels like a rare scenario.
+            && blContext.lookup().onlyMappedAsRuntimeField(name())) {
+            var reader = new NumberType.NumberFallbackSyntheticSourceReader(NumberType.LONG, null, true) {
+                @Override
+                public void writeToBlock(List<Number> values, BlockLoader.Builder blockBuilder) {
+                    var builder = (BlockLoader.LongBuilder) blockBuilder;
+                    for (var value : values) {
+                        builder.appendLong(value.longValue());
+                    }
+                }
+            };
+
+            return new FallbackSyntheticSourceBlockLoader(
+                reader,
+                name(),
+                IgnoredSourceFieldMapper.ignoredSourceFormat(indexSettings.getIndexVersionCreated())
+            ) {
+                @Override
+                public Builder builder(BlockFactory factory, int expectedCount) {
+                    return factory.longs(expectedCount);
+                }
+            };
+        } else {
+            return new LongScriptBlockDocValuesReader.LongScriptBlockLoader(leafFactory(blContext.lookup()));
+        }
     }
 
     @Override

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

@@ -534,6 +534,7 @@ public void setLookupProviders(
         // as well as runtime fields loaded from _source that do need a source provider as part of executing the query
         this.lookup = new SearchLookup(
             this::getFieldType,
+            fieldName -> mappingLookup.getMapper(fieldName) == null,
             (fieldType, searchLookup, fielddataOperation) -> indexFieldDataLookup.apply(
                 fieldType,
                 new FieldDataContext(

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

@@ -46,6 +46,7 @@ public class SearchLookup implements SourceProvider {
     private final Set<String> fieldChain;
     private final SourceProvider sourceProvider;
     private final Function<String, MappedFieldType> fieldTypeLookup;
+    private final Function<String, Boolean> onlyMappedAsRuntimeField;
     private final TriFunction<
         MappedFieldType,
         Supplier<SearchLookup>,
@@ -64,7 +65,7 @@ public SearchLookup(
         TriFunction<MappedFieldType, Supplier<SearchLookup>, MappedFieldType.FielddataOperation, IndexFieldData<?>> fieldDataLookup,
         SourceProvider sourceProvider
     ) {
-        this(fieldTypeLookup, fieldDataLookup, sourceProvider, LeafFieldLookupProvider.fromStoredFields());
+        this(fieldTypeLookup, fieldName -> false, fieldDataLookup, sourceProvider, LeafFieldLookupProvider.fromStoredFields());
     }
 
     /**
@@ -76,11 +77,13 @@ public SearchLookup(
      */
     public SearchLookup(
         Function<String, MappedFieldType> fieldTypeLookup,
+        Function<String, Boolean> onlyMappedAsRuntimeField,
         TriFunction<MappedFieldType, Supplier<SearchLookup>, MappedFieldType.FielddataOperation, IndexFieldData<?>> fieldDataLookup,
         SourceProvider sourceProvider,
         Function<LeafReaderContext, LeafFieldLookupProvider> fieldLookupProvider
     ) {
         this.fieldTypeLookup = fieldTypeLookup;
+        this.onlyMappedAsRuntimeField = onlyMappedAsRuntimeField;
         this.fieldChain = Collections.emptySet();
         this.sourceProvider = sourceProvider;
         this.fieldDataLookup = fieldDataLookup;
@@ -98,6 +101,7 @@ private SearchLookup(SearchLookup searchLookup, Set<String> fieldChain) {
         this.fieldChain = Collections.unmodifiableSet(fieldChain);
         this.sourceProvider = searchLookup.sourceProvider;
         this.fieldTypeLookup = searchLookup.fieldTypeLookup;
+        this.onlyMappedAsRuntimeField = searchLookup.onlyMappedAsRuntimeField;
         this.fieldDataLookup = searchLookup.fieldDataLookup;
         this.fieldLookupProvider = searchLookup.fieldLookupProvider;
     }
@@ -106,6 +110,7 @@ private SearchLookup(SearchLookup searchLookup, SourceProvider sourceProvider) {
         this.fieldChain = searchLookup.fieldChain;
         this.sourceProvider = sourceProvider;
         this.fieldTypeLookup = searchLookup.fieldTypeLookup;
+        this.onlyMappedAsRuntimeField = searchLookup.onlyMappedAsRuntimeField;
         this.fieldDataLookup = searchLookup.fieldDataLookup;
         this.fieldLookupProvider = searchLookup.fieldLookupProvider;
     }
@@ -144,6 +149,13 @@ public MappedFieldType fieldType(String fieldName) {
         return fieldTypeLookup.apply(fieldName);
     }
 
+    /**
+     * @return whether a field is only mapped as runtime field. A runtime and normal field can share the same name.
+     */
+    public boolean onlyMappedAsRuntimeField(String fieldName) {
+        return onlyMappedAsRuntimeField.apply(fieldName);
+    }
+
     public IndexFieldData<?> getForField(MappedFieldType fieldType, MappedFieldType.FielddataOperation options) {
         return fieldDataLookup.apply(fieldType, () -> forkAndTrackFieldReferences(fieldType.name()), options);
     }

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

@@ -29,8 +29,11 @@
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.geo.ShapeRelation;
 import org.elasticsearch.common.lucene.search.function.ScriptScoreQuery;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.fielddata.LongScriptFieldData;
 import org.elasticsearch.index.fielddata.ScriptDocValues;
@@ -42,6 +45,10 @@
 import org.elasticsearch.script.ScriptFactory;
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.search.MultiValueMode;
+import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.xcontent.XContentFactory;
+import org.elasticsearch.xcontent.XContentParserConfiguration;
+import org.elasticsearch.xcontent.XContentType;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -52,6 +59,8 @@
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.nullValue;
 
 public class LongScriptFieldTypeTests extends AbstractNonTextScriptFieldTypeTestCase {
 
@@ -302,6 +311,83 @@ public void testBlockLoader() throws IOException {
         }
     }
 
+    public void testBlockLoaderSourceOnlyRuntimeField() throws IOException {
+        try (
+            Directory directory = newDirectory();
+            RandomIndexWriter iw = new RandomIndexWriter(random(), directory, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))
+        ) {
+            iw.addDocuments(
+                List.of(
+                    List.of(new StoredField("_source", new BytesRef("{\"test\": [1]}"))),
+                    List.of(new StoredField("_source", new BytesRef("{\"test\": [2]}")))
+                )
+            );
+            try (DirectoryReader reader = iw.getReader()) {
+                LongScriptFieldType fieldType = simpleSourceOnlyMappedFieldType();
+
+                // Assert implementations:
+                BlockLoader loader = fieldType.blockLoader(blContext(Settings.EMPTY, true));
+                assertThat(loader, instanceOf(LongScriptBlockDocValuesReader.LongScriptBlockLoader.class));
+                // ignored source doesn't support column at a time loading:
+                var columnAtATimeLoader = loader.columnAtATimeReader(reader.leaves().getFirst());
+                assertThat(columnAtATimeLoader, instanceOf(LongScriptBlockDocValuesReader.class));
+                var rowStrideReader = loader.rowStrideReader(reader.leaves().getFirst());
+                assertThat(rowStrideReader, instanceOf(LongScriptBlockDocValuesReader.class));
+
+                // Assert values:
+                assertThat(blockLoaderReadValuesFromColumnAtATimeReader(reader, fieldType, 0), equalTo(List.of(1L, 2L)));
+                assertThat(blockLoaderReadValuesFromColumnAtATimeReader(reader, fieldType, 1), equalTo(List.of(2L)));
+                assertThat(blockLoaderReadValuesFromRowStrideReader(reader, fieldType), equalTo(List.of(1L, 2L)));
+            }
+        }
+    }
+
+    public void testBlockLoaderSourceOnlyRuntimeFieldWithSyntheticSource() throws IOException {
+        var settings = Settings.builder().put("index.mapping.source.mode", "synthetic").build();
+        try (
+            Directory directory = newDirectory();
+            RandomIndexWriter iw = new RandomIndexWriter(random(), directory, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))
+        ) {
+
+            var document1 = createDocumentWithIgnoredSource("[1]");
+            var document2 = createDocumentWithIgnoredSource("[2]");
+
+            iw.addDocuments(List.of(document1, document2));
+            try (DirectoryReader reader = iw.getReader()) {
+                LongScriptFieldType fieldType = simpleSourceOnlyMappedFieldType();
+
+                // Assert implementations:
+                BlockLoader loader = fieldType.blockLoader(blContext(settings, true));
+                assertThat(loader, instanceOf(FallbackSyntheticSourceBlockLoader.class));
+                // ignored source doesn't support column at a time loading:
+                var columnAtATimeLoader = loader.columnAtATimeReader(reader.leaves().getFirst());
+                assertThat(columnAtATimeLoader, nullValue());
+                var rowStrideReader = loader.rowStrideReader(reader.leaves().getFirst());
+                assertThat(
+                    rowStrideReader.getClass().getName(),
+                    equalTo("org.elasticsearch.index.mapper.FallbackSyntheticSourceBlockLoader$IgnoredSourceRowStrideReader")
+                );
+
+                // Assert values:
+                assertThat(blockLoaderReadValuesFromRowStrideReader(settings, reader, fieldType, true), equalTo(List.of(1L, 2L)));
+            }
+        }
+    }
+
+    private static LuceneDocument createDocumentWithIgnoredSource(String bytes) throws IOException {
+        var doc = new LuceneDocument();
+        var parser = XContentHelper.createParser(
+            XContentParserConfiguration.EMPTY,
+            new BytesArray(bytes),
+            XContentFactory.xContent(XContentType.JSON).type()
+        );
+        parser.nextToken();
+        var nameValue = new IgnoredSourceFieldMapper.NameValue("test", 0, XContentDataHelper.encodeToken(parser), doc);
+        var ignoredSourceFormat = IgnoredSourceFieldMapper.ignoredSourceFormat(IndexVersion.current());
+        ignoredSourceFormat.writeIgnoredFields(List.of(nameValue));
+        return doc;
+    }
+
     @Override
     protected Query randomTermsQuery(MappedFieldType ft, SearchExecutionContext ctx) {
         return ft.termsQuery(List.of(randomLong()), ctx);
@@ -312,6 +398,10 @@ protected LongScriptFieldType simpleMappedFieldType() {
         return build("read_foo", Map.of(), OnScriptError.FAIL);
     }
 
+    private LongScriptFieldType simpleSourceOnlyMappedFieldType() {
+        return build("read_test", Map.of(), OnScriptError.FAIL);
+    }
+
     @Override
     protected LongScriptFieldType loopFieldType() {
         return build("loop", Map.of(), OnScriptError.FAIL);
@@ -329,6 +419,32 @@ protected LongScriptFieldType build(String code, Map<String, Object> params, OnS
 
     private static LongFieldScript.Factory factory(Script script) {
         switch (script.getIdOrCode()) {
+            case "read_test":
+                return new LongFieldScript.Factory() {
+                    @Override
+                    public LongFieldScript.LeafFactory newFactory(
+                        String fieldName,
+                        Map<String, Object> params,
+                        SearchLookup lookup,
+                        OnScriptError onScriptError
+                    ) {
+                        return (ctx) -> new LongFieldScript(fieldName, params, lookup, onScriptError, ctx) {
+                            @Override
+                            @SuppressWarnings("unchecked")
+                            public void execute() {
+                                Map<String, Object> source = (Map<String, Object>) this.getParams().get("_source");
+                                for (Object foo : (List<?>) source.get("test")) {
+                                    emit(((Number) foo).longValue());
+                                }
+                            };
+                        };
+                    }
+
+                    @Override
+                    public boolean isParsedFromSource() {
+                        return true;
+                    }
+                };
             case "read_foo":
                 return (fieldName, params, lookup, onScriptError) -> (ctx) -> new LongFieldScript(
                     fieldName,