Improve block loader for source only runtime fields of type keyword (elastic#135026)

jordan-powers · web-flow · commit 0022605a02df · 2025-09-19T09:18:53.000-07:00
This patch updates source-only keyword scripts to use the FallbackSyntheticSourceBlockLoader instead of generic KeywordScriptBlockLoader to retrieve the values out of ignored_source when synthetic source is enabled. Relates to elastic#134121
diff --git a/docs/changelog/135026.yaml b/docs/changelog/135026.yaml
@@ -0,0 +1,5 @@
+pr: 135026
+summary: Improve block loader for source only runtime fields of type keyword
+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
@@ -40,6 +40,7 @@
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 import static org.elasticsearch.search.SearchService.ALLOW_EXPENSIVE_QUERIES;
 
@@ -205,15 +206,32 @@ protected final LeafFactory leafFactory(SearchLookup searchLookup) {
         }
     }
 
+    protected final FallbackSyntheticSourceBlockLoader numericFallbackSyntheticSourceBlockLoader(
+        BlockLoaderContext blContext,
+        NumberFieldMapper.NumberType numberType,
+        BiFunction<BlockLoader.BlockFactory, Integer, BlockLoader.Builder> builderSupplier,
+        BiConsumer<List<Number>, BlockLoader.Builder> writeToBlock
+    ) {
+        return fallbackSyntheticSourceBlockLoader(
+            blContext,
+            builderSupplier,
+            () -> new NumberFieldMapper.NumberType.NumberFallbackSyntheticSourceReader(numberType, null, true) {
+                @Override
+                public void writeToBlock(List<Number> values, BlockLoader.Builder blockBuilder) {
+                    writeToBlock.accept(values, blockBuilder);
+                }
+            }
+        );
+    }
+
     /**
      * Returns synthetic source fallback block loader if source mode is synthetic, runtime field is source only and field is only mapped
      * as a runtime field.
      */
     protected final FallbackSyntheticSourceBlockLoader fallbackSyntheticSourceBlockLoader(
         BlockLoaderContext blContext,
-        NumberFieldMapper.NumberType numberType,
         BiFunction<BlockLoader.BlockFactory, Integer, BlockLoader.Builder> builderSupplier,
-        BiConsumer<List<Number>, BlockLoader.Builder> writeToBlock
+        Supplier<FallbackSyntheticSourceBlockLoader.Reader<?>> readerSupplier
     ) {
         var indexSettings = blContext.indexSettings();
         // A runtime and normal field can share the same name.
@@ -222,12 +240,7 @@ protected final FallbackSyntheticSourceBlockLoader fallbackSyntheticSourceBlockL
         if (isParsedFromSource
             && indexSettings.getIndexMappingSourceMode() == SourceFieldMapper.Mode.SYNTHETIC
             && blContext.lookup().onlyMappedAsRuntimeField(name())) {
-            var reader = new NumberFieldMapper.NumberType.NumberFallbackSyntheticSourceReader(numberType, null, true) {
-                @Override
-                public void writeToBlock(List<Number> values, BlockLoader.Builder blockBuilder) {
-                    writeToBlock.accept(values, blockBuilder);
-                }
-            };
+            var reader = readerSupplier.get();
 
             return new FallbackSyntheticSourceBlockLoader(
                 reader,
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/DoubleScriptFieldType.java b/server/src/main/java/org/elasticsearch/index/mapper/DoubleScriptFieldType.java
@@ -109,7 +109,7 @@ public DocValueFormat docValueFormat(String format, ZoneId timeZone) {
 
     @Override
     public BlockLoader blockLoader(BlockLoaderContext blContext) {
-        var fallbackSyntheticSourceBlockLoader = fallbackSyntheticSourceBlockLoader(
+        var fallbackSyntheticSourceBlockLoader = numericFallbackSyntheticSourceBlockLoader(
             blContext,
             NumberType.DOUBLE,
             BlockLoader.BlockFactory::doubles,
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/KeywordScriptFieldType.java b/server/src/main/java/org/elasticsearch/index/mapper/KeywordScriptFieldType.java
@@ -33,9 +33,12 @@
 import org.elasticsearch.search.runtime.StringScriptFieldTermQuery;
 import org.elasticsearch.search.runtime.StringScriptFieldTermsQuery;
 import org.elasticsearch.search.runtime.StringScriptFieldWildcardQuery;
+import org.elasticsearch.xcontent.XContentParser;
 
+import java.io.IOException;
 import java.time.ZoneId;
 import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -112,9 +115,44 @@ public Object valueForDisplay(Object value) {
         return binaryValue.utf8ToString();
     }
 
+    private FallbackSyntheticSourceBlockLoader.Reader<?> fallbackSyntheticSourceBlockLoaderReader() {
+        return new FallbackSyntheticSourceBlockLoader.SingleValueReader<BytesRef>(null) {
+            @Override
+            public void convertValue(Object value, List<BytesRef> accumulator) {
+                accumulator.add((BytesRef) value);
+            }
+
+            @Override
+            public void parseNonNullValue(XContentParser parser, List<BytesRef> accumulator) throws IOException {
+                assert parser.currentToken() == XContentParser.Token.VALUE_STRING : "Unexpected token " + parser.currentToken();
+
+                var utfBytes = parser.optimizedText().bytes();
+                accumulator.add(new BytesRef(utfBytes.bytes(), utfBytes.offset(), utfBytes.length()));
+            }
+
+            @Override
+            public void writeToBlock(List<BytesRef> values, BlockLoader.Builder blockBuilder) {
+                var bytesRefBuilder = (BlockLoader.BytesRefBuilder) blockBuilder;
+
+                for (var value : values) {
+                    bytesRefBuilder.appendBytesRef(value);
+                }
+            }
+        };
+    }
+
     @Override
     public BlockLoader blockLoader(BlockLoaderContext blContext) {
-        return new KeywordScriptBlockDocValuesReader.KeywordScriptBlockLoader(leafFactory(blContext.lookup()));
+        var fallbackSyntheticSourceBlockLoader = fallbackSyntheticSourceBlockLoader(
+            blContext,
+            BlockLoader.BlockFactory::bytesRefs,
+            this::fallbackSyntheticSourceBlockLoaderReader
+        );
+        if (fallbackSyntheticSourceBlockLoader != null) {
+            return fallbackSyntheticSourceBlockLoader;
+        } else {
+            return new KeywordScriptBlockDocValuesReader.KeywordScriptBlockLoader(leafFactory(blContext.lookup()));
+        }
     }
 
     @Override
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/LongScriptFieldType.java b/server/src/main/java/org/elasticsearch/index/mapper/LongScriptFieldType.java
@@ -109,7 +109,7 @@ public DocValueFormat docValueFormat(String format, ZoneId timeZone) {
 
     @Override
     public BlockLoader blockLoader(BlockLoaderContext blContext) {
-        var fallbackSyntheticSourceBlockLoader = fallbackSyntheticSourceBlockLoader(
+        var fallbackSyntheticSourceBlockLoader = numericFallbackSyntheticSourceBlockLoader(
             blContext,
             NumberType.LONG,
             BlockLoader.BlockFactory::longs,
diff --git a/server/src/test/java/org/elasticsearch/index/mapper/DoubleScriptFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/DoubleScriptFieldTypeTests.java
@@ -49,7 +49,6 @@
 import java.util.Map;
 
 import static java.util.Collections.emptyMap;
-import static org.elasticsearch.index.mapper.LongScriptFieldTypeTests.createDocumentWithIgnoredSource;
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
diff --git a/server/src/test/java/org/elasticsearch/index/mapper/KeywordScriptFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/KeywordScriptFieldTypeTests.java
@@ -29,6 +29,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.automaton.Operations;
 import org.elasticsearch.common.lucene.search.function.ScriptScoreQuery;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.Fuzziness;
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.fielddata.BinaryScriptFieldData;
@@ -44,6 +45,7 @@
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.script.StringFieldScript;
 import org.elasticsearch.search.MultiValueMode;
+import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -53,6 +55,8 @@
 import static java.util.Collections.emptyMap;
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.nullValue;
 
 public class KeywordScriptFieldTypeTests extends AbstractScriptFieldTypeTestCase {
 
@@ -421,6 +425,105 @@ 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\": [\"cat\"]}"))),
+                    List.of(new StoredField("_source", new BytesRef("{\"test\": [\"dog\"]}")))
+                )
+            );
+            try (DirectoryReader reader = iw.getReader()) {
+                KeywordScriptFieldType fieldType = simpleSourceOnlyMappedFieldType();
+
+                // Assert implementations:
+                BlockLoader loader = fieldType.blockLoader(blContext(Settings.EMPTY, true));
+                assertThat(loader, instanceOf(KeywordScriptBlockDocValuesReader.KeywordScriptBlockLoader.class));
+                // ignored source doesn't support column at a time loading:
+                var columnAtATimeLoader = loader.columnAtATimeReader(reader.leaves().getFirst());
+                assertThat(columnAtATimeLoader, instanceOf(KeywordScriptBlockDocValuesReader.class));
+                var rowStrideReader = loader.rowStrideReader(reader.leaves().getFirst());
+                assertThat(rowStrideReader, instanceOf(KeywordScriptBlockDocValuesReader.class));
+
+                var catBytes = new BytesRef("cat");
+                var dogBytes = new BytesRef("dog");
+
+                // Assert values:
+                assertThat(blockLoaderReadValuesFromColumnAtATimeReader(reader, fieldType, 0), equalTo(List.of(catBytes, dogBytes)));
+                assertThat(blockLoaderReadValuesFromColumnAtATimeReader(reader, fieldType, 1), equalTo(List.of(dogBytes)));
+                assertThat(blockLoaderReadValuesFromRowStrideReader(reader, fieldType), equalTo(List.of(catBytes, dogBytes)));
+            }
+        }
+    }
+
+    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("[\"cat\"]");
+            var document2 = createDocumentWithIgnoredSource("[\"dog\"]");
+
+            iw.addDocuments(List.of(document1, document2));
+            try (DirectoryReader reader = iw.getReader()) {
+                KeywordScriptFieldType 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(new BytesRef("cat"), new BytesRef("dog")))
+                );
+            }
+        }
+    }
+
+    private KeywordScriptFieldType simpleSourceOnlyMappedFieldType() {
+        Script script = new Script(ScriptType.INLINE, "test", "", emptyMap());
+        StringFieldScript.Factory factory = new StringFieldScript.Factory() {
+            @Override
+            public StringFieldScript.LeafFactory newFactory(
+                String fieldName,
+                Map<String, Object> params,
+                SearchLookup searchLookup,
+                OnScriptError onScriptError
+            ) {
+                return ctx -> new StringFieldScript(fieldName, params, searchLookup, 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((String) foo);
+                        }
+                    }
+                };
+            }
+
+            @Override
+            public boolean isParsedFromSource() {
+                return true;
+            }
+        };
+        return new KeywordScriptFieldType("test", factory, script, emptyMap(), OnScriptError.FAIL);
+    }
+
     @Override
     protected KeywordScriptFieldType simpleMappedFieldType() {
         return build("read_foo", Map.of(), OnScriptError.FAIL);
diff --git a/server/src/test/java/org/elasticsearch/index/mapper/LongScriptFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/LongScriptFieldTypeTests.java
@@ -29,11 +29,9 @@
 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;
@@ -46,9 +44,6 @@
 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;
@@ -374,20 +369,6 @@ public void testBlockLoaderSourceOnlyRuntimeFieldWithSyntheticSource() throws IO
         }
     }
 
-    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);
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
@@ -24,6 +24,7 @@
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.geo.ShapeRelation;
 import org.elasticsearch.common.settings.Settings;
@@ -46,6 +47,8 @@
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentFactory;
+import org.elasticsearch.xcontent.XContentParserConfiguration;
+import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xcontent.json.JsonXContent;
 
 import java.io.IOException;
@@ -621,4 +624,18 @@ protected <T> T compileScript(Script script, ScriptContext<T> context) {
         boolean deterministicSource = "deterministic_source".equals(script.getIdOrCode());
         return deterministicSource ? (T) parseFromSource() : (T) dummyScript();
     }
+
+    protected 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;
+    }
 }
