diff --git a/docs/changelog/134117.yaml b/docs/changelog/134117.yaml new file mode 100644 index 0000000000000..fd318442a2b2a --- /dev/null +++ b/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: [] 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 1140a45fe0308..7f4b287018bed 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/AbstractScriptFieldType.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/AbstractScriptFieldType.java @@ -49,7 +49,7 @@ public abstract class AbstractScriptFieldType extends MappedFieldTy protected final Script script; private final Function factory; private final boolean isResultDeterministic; - private final boolean isParsedFromSource; + protected final boolean isParsedFromSource; protected AbstractScriptFieldType( String name, diff --git a/server/src/main/java/org/elasticsearch/index/mapper/DocumentParser.java b/server/src/main/java/org/elasticsearch/index/mapper/DocumentParser.java index 87017a24765dc..10e53a16ed076 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/DocumentParser.java +++ b/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 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 9cbf60dd87240..cbfd9ad2c3570 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/LongScriptFieldType.java +++ b/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 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 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 6e0bf8414d84d..c644c146569cd 100644 --- a/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java +++ b/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java @@ -528,6 +528,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( 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 4e932ea47ca7e..fff4eec6d8b3d 100644 --- a/server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java +++ b/server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java @@ -46,6 +46,7 @@ public class SearchLookup implements SourceProvider { private final Set fieldChain; private final SourceProvider sourceProvider; private final Function fieldTypeLookup; + private final Function onlyMappedAsRuntimeField; private final TriFunction< MappedFieldType, Supplier, @@ -64,7 +65,7 @@ public SearchLookup( TriFunction, 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 fieldTypeLookup, + Function onlyMappedAsRuntimeField, TriFunction, MappedFieldType.FielddataOperation, IndexFieldData> fieldDataLookup, SourceProvider sourceProvider, Function 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 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); } diff --git a/server/src/test/java/org/elasticsearch/index/mapper/LongScriptFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/LongScriptFieldTypeTests.java index 01f96a1a4b1be..67b0b47adf4bd 100644 --- a/server/src/test/java/org/elasticsearch/index/mapper/LongScriptFieldTypeTests.java +++ b/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 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 params, + SearchLookup lookup, + OnScriptError onScriptError + ) { + return (ctx) -> new LongFieldScript(fieldName, params, lookup, onScriptError, ctx) { + @Override + @SuppressWarnings("unchecked") + public void execute() { + Map source = (Map) 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, diff --git a/server/src/test/java/org/elasticsearch/index/mapper/RuntimeFieldSourceProviderOptimizationTests.java b/server/src/test/java/org/elasticsearch/index/mapper/RuntimeFieldSourceProviderOptimizationTests.java index 03234f12e421c..95da43177406c 100644 --- a/server/src/test/java/org/elasticsearch/index/mapper/RuntimeFieldSourceProviderOptimizationTests.java +++ b/server/src/test/java/org/elasticsearch/index/mapper/RuntimeFieldSourceProviderOptimizationTests.java @@ -9,18 +9,24 @@ package org.elasticsearch.index.mapper; +import org.apache.lucene.document.Document; 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.cluster.metadata.IndexMetadata; import org.elasticsearch.common.bytes.BytesArray; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.index.IndexSettings; +import org.elasticsearch.index.IndexVersion; +import org.elasticsearch.index.fielddata.LongScriptDocValues; +import org.elasticsearch.index.fielddata.LongScriptFieldData; import org.elasticsearch.script.LongFieldScript; import org.elasticsearch.search.lookup.SearchLookup; import org.elasticsearch.test.ESSingleNodeTestCase; +import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xcontent.XContentType; import java.io.IOException; @@ -30,6 +36,8 @@ import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.notNullValue; /** * Tests that source provider optimization that filters _source based on the same of source only runtime fields kick in. @@ -41,7 +49,8 @@ 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); + var settings = Settings.builder().put("index.mapping.source.mode", "synthetic").build(); + var indexService = createIndex("test-index", settings, mapping); int numDocs = 256; try (Directory directory = newDirectory(); IndexWriter iw = new IndexWriter(directory, new IndexWriterConfig())) { @@ -77,12 +86,15 @@ public void testWithSourceProviderOptimization() throws IOException { 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)); + // Test that script runtime field data works as expected with the optimization: + var fieldData = (LongScriptFieldData) context.getForField(fieldType, MappedFieldType.FielddataOperation.SCRIPT); + var leafFieldData = fieldData.load(leafReaderContext); + var sortedNumericDocValues = (LongScriptDocValues) leafFieldData.getLongValues(); + for (int i = 0; i < 256; i++) { + boolean result = sortedNumericDocValues.advanceExact(i); + assertThat(result, equalTo(true)); + assertThat(sortedNumericDocValues.docValueCount(), equalTo(1)); + assertThat(sortedNumericDocValues.nextValue(), equalTo((long) i)); } } } @@ -125,27 +137,66 @@ public void testWithoutSourceProviderOptimization() throws IOException { 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)); + // Test that script runtime field data works as expected with the optimization: + var fieldData = (LongScriptFieldData) context.getForField(fieldType, MappedFieldType.FielddataOperation.SCRIPT); + var leafFieldData = fieldData.load(leafReaderContext); + var sortedNumericDocValues = (LongScriptDocValues) leafFieldData.getLongValues(); + for (int i = 0; i < 256; i++) { + boolean result = sortedNumericDocValues.advanceExact(i); + assertThat(result, equalTo(true)); + assertThat(sortedNumericDocValues.docValueCount(), equalTo(1)); + assertThat(sortedNumericDocValues.nextValue(), equalTo((long) i)); } } } } - static MappedFieldType.BlockLoaderContext blContext(SearchLookup lookup) { + public void testNormalAndRuntimeFieldWithSameName() throws IOException { + var mapping = jsonBuilder().startObject().startObject("runtime"); + mapping.startObject("field").field("type", "long").endObject(); + mapping.startObject("field2").field("type", "long").endObject(); + mapping.endObject().startObject("properties"); + mapping.startObject("field").field("type", "long").endObject(); + mapping.endObject().endObject(); + + var settings = Settings.builder().put("index.mapping.source.mode", "synthetic").build(); + var indexService = createIndex("test-index", settings, mapping); + var fieldType1 = indexService.mapperService().fieldType("field"); + assertThat(fieldType1, notNullValue()); + var fieldType2 = indexService.mapperService().fieldType("field2"); + assertThat(fieldType2, notNullValue()); + + // Assert implementations: + try (Directory directory = newDirectory(); IndexWriter iw = new IndexWriter(directory, new IndexWriterConfig())) { + iw.addDocument(new Document()); + try (var indexReader = DirectoryReader.open(iw)) { + var searcher = new IndexSearcher(indexReader); + var context = indexService.newSearchExecutionContext(0, 0, searcher, () -> 1L, null, Map.of()); + + // field name 'field' is both mapped as runtime and normal field and so LongScriptBlockLoader is expected: + BlockLoader loader = fieldType1.blockLoader(blContext(settings, context.lookup())); + assertThat(loader, instanceOf(LongScriptBlockDocValuesReader.LongScriptBlockLoader.class)); + + // field name 'field2' is just mapped as runtime field and so FallbackSyntheticSourceBlockLoader is expected: + BlockLoader loader2 = fieldType2.blockLoader(blContext(settings, context.lookup())); + assertThat(loader2, instanceOf(FallbackSyntheticSourceBlockLoader.class)); + } + } + + } + + static MappedFieldType.BlockLoaderContext blContext(Settings settings, SearchLookup lookup) { + String indexName = "test_index"; + var imd = IndexMetadata.builder(indexName).settings(ESTestCase.indexSettings(IndexVersion.current(), 1, 1).put(settings)).build(); return new MappedFieldType.BlockLoaderContext() { @Override public String indexName() { - throw new UnsupportedOperationException(); + return indexName; } @Override public IndexSettings indexSettings() { - throw new UnsupportedOperationException(); + return new IndexSettings(imd, settings); } @Override @@ -174,4 +225,5 @@ public FieldNamesFieldMapper.FieldNamesFieldType fieldNames() { } }; } + } 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 f679d47274b46..de53d080fe398 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 @@ -22,21 +22,27 @@ import org.apache.lucene.tests.index.RandomIndexWriter; import org.apache.lucene.util.BytesRef; import org.elasticsearch.ElasticsearchException; +import org.elasticsearch.cluster.metadata.IndexMetadata; import org.elasticsearch.common.Strings; import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.geo.ShapeRelation; +import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.xcontent.XContentHelper; import org.elasticsearch.index.IndexSettings; +import org.elasticsearch.index.IndexVersion; import org.elasticsearch.index.fielddata.FieldDataContext; import org.elasticsearch.index.fielddata.IndexFieldDataCache; +import org.elasticsearch.index.fieldvisitor.StoredFieldLoader; import org.elasticsearch.index.query.ExistsQueryBuilder; import org.elasticsearch.index.query.SearchExecutionContext; import org.elasticsearch.indices.breaker.NoneCircuitBreakerService; import org.elasticsearch.script.Script; import org.elasticsearch.script.ScriptContext; import org.elasticsearch.script.ScriptFactory; +import org.elasticsearch.search.lookup.LeafFieldLookupProvider; import org.elasticsearch.search.lookup.SearchLookup; import org.elasticsearch.search.lookup.SourceProvider; +import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xcontent.ToXContent; import org.elasticsearch.xcontent.XContentBuilder; import org.elasticsearch.xcontent.XContentFactory; @@ -292,6 +298,7 @@ protected boolean supportsRangeQueries() { protected static SearchExecutionContext mockContext(boolean allowExpensiveQueries, MappedFieldType mappedFieldType) { return mockContext( allowExpensiveQueries, + false, mappedFieldType, SourceProvider.fromLookup(MappingLookup.EMPTY, null, SourceFieldMetrics.NOOP) ); @@ -299,6 +306,7 @@ protected static SearchExecutionContext mockContext(boolean allowExpensiveQuerie protected static SearchExecutionContext mockContext( boolean allowExpensiveQueries, + boolean fieldOnlyMappedAsRuntimeField, MappedFieldType mappedFieldType, SourceProvider sourceProvider ) { @@ -309,9 +317,11 @@ protected static SearchExecutionContext mockContext( when(context.allowExpensiveQueries()).thenReturn(allowExpensiveQueries); SearchLookup lookup = new SearchLookup( context::getFieldType, + (fieldName) -> fieldOnlyMappedAsRuntimeField, (mft, lookupSupplier, fdo) -> mft.fielddataBuilder(new FieldDataContext("test", null, lookupSupplier, context::sourcePath, fdo)) .build(null, null), - sourceProvider + sourceProvider, + LeafFieldLookupProvider.fromStoredFields() ); when(context.lookup()).thenReturn(lookup); when(context.getForField(any(), any())).then(args -> { @@ -465,7 +475,16 @@ public final void testIsParsedFromSource() throws IOException { protected final List blockLoaderReadValuesFromColumnAtATimeReader(DirectoryReader reader, MappedFieldType fieldType, int offset) throws IOException { - BlockLoader loader = fieldType.blockLoader(blContext()); + return blockLoaderReadValuesFromColumnAtATimeReader(Settings.EMPTY, reader, fieldType, offset); + } + + protected final List blockLoaderReadValuesFromColumnAtATimeReader( + Settings settings, + DirectoryReader reader, + MappedFieldType fieldType, + int offset + ) throws IOException { + BlockLoader loader = fieldType.blockLoader(blContext(settings, true)); List all = new ArrayList<>(); for (LeafReaderContext ctx : reader.leaves()) { TestBlock block = (TestBlock) loader.columnAtATimeReader(ctx).read(TestBlock.factory(), TestBlock.docs(ctx), offset, false); @@ -478,13 +497,29 @@ protected final List blockLoaderReadValuesFromColumnAtATimeReader(Direct protected final List blockLoaderReadValuesFromRowStrideReader(DirectoryReader reader, MappedFieldType fieldType) throws IOException { - BlockLoader loader = fieldType.blockLoader(blContext()); + return blockLoaderReadValuesFromRowStrideReader(Settings.EMPTY, reader, fieldType, false); + } + + protected final List blockLoaderReadValuesFromRowStrideReader( + Settings settings, + DirectoryReader reader, + MappedFieldType fieldType, + boolean fieldOnlyMappedAsRuntimeField + ) throws IOException { + BlockLoader loader = fieldType.blockLoader(blContext(settings, fieldOnlyMappedAsRuntimeField)); List all = new ArrayList<>(); for (LeafReaderContext ctx : reader.leaves()) { BlockLoader.RowStrideReader blockReader = loader.rowStrideReader(ctx); BlockLoader.Builder builder = loader.builder(TestBlock.factory(), ctx.reader().numDocs()); + + assert loader.rowStrideStoredFieldSpec().requiresSource() == false; + BlockLoaderStoredFieldsFromLeafLoader storedFields = new BlockLoaderStoredFieldsFromLeafLoader( + StoredFieldLoader.fromSpec(loader.rowStrideStoredFieldSpec()).getLoader(ctx, null), + null + ); for (int i = 0; i < ctx.reader().numDocs(); i++) { - blockReader.read(i, null, builder); + storedFields.advanceTo(i); + blockReader.read(i, storedFields, builder); } TestBlock block = (TestBlock) builder.build(); for (int i = 0; i < block.size(); i++) { @@ -494,16 +529,18 @@ protected final List blockLoaderReadValuesFromRowStrideReader(DirectoryR return all; } - private MappedFieldType.BlockLoaderContext blContext() { + protected MappedFieldType.BlockLoaderContext blContext(Settings settings, boolean fieldOnlyMappedAsRuntimeField) { + String indexName = "test_index"; + var imd = IndexMetadata.builder(indexName).settings(ESTestCase.indexSettings(IndexVersion.current(), 1, 1).put(settings)).build(); return new MappedFieldType.BlockLoaderContext() { @Override public String indexName() { - throw new UnsupportedOperationException(); + return indexName; } @Override public IndexSettings indexSettings() { - throw new UnsupportedOperationException(); + return new IndexSettings(imd, settings); } @Override @@ -513,7 +550,12 @@ public MappedFieldType.FieldExtractPreference fieldExtractPreference() { @Override public SearchLookup lookup() { - return mockContext().lookup(); + return mockContext( + true, + fieldOnlyMappedAsRuntimeField, + null, + SourceProvider.fromLookup(MappingLookup.EMPTY, null, SourceFieldMetrics.NOOP) + ).lookup(); } @Override