Skip to content

Commit f967b21

Browse files
lktslkts
committed
Introduce BlockLoaderTestCase
1 parent bf3edff commit f967b21

File tree

3 files changed

+237
-3
lines changed

3 files changed

+237
-3
lines changed

‎server/src/test/java/org/elasticsearch/index/mapper/blockloader/KeywordFieldBlockLoaderTests.java‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.index.mapper.blockloader;
11+
12+
import org.apache.lucene.util.BytesRef;
13+
import org.elasticsearch.index.mapper.BlockLoaderTestCase;
14+
import org.elasticsearch.logsdb.datageneration.FieldType;
15+
16+
import java.util.List;
17+
import java.util.Map;
18+
import java.util.Objects;
19+
import java.util.stream.Collectors;
20+
21+
public class KeywordFieldBlockLoaderTests extends BlockLoaderTestCase {
22+
public KeywordFieldBlockLoaderTests() {
23+
super(FieldType.KEYWORD);
24+
}
25+
26+
@SuppressWarnings("unchecked")
27+
@Override
28+
protected Object expected(Map<String, Object> fieldMapping, Object value, boolean syntheticSource) {
29+
if (value == null) {
30+
return null;
31+
}
32+
if (value instanceof String s) {
33+
return convert(s);
34+
}
35+
36+
var nonNullStream = ((List<String>) value).stream().filter(Objects::nonNull);
37+
38+
if ((boolean) fieldMapping.getOrDefault("doc_values", false)) {
39+
// Sorted and no duplicates
40+
return maybeFoldList(nonNullStream.collect(Collectors.toSet()).stream().sorted().map(this::convert).toList());
41+
}
42+
43+
if ((boolean) fieldMapping.getOrDefault("store", false)) {
44+
return maybeFoldList(nonNullStream.map(this::convert).toList());
45+
}
46+
47+
// Using source (either stored or synthetic).
48+
// Original order is preserved and values longer than ignore_above are returned.
49+
// TODO actual ignore_above support in data generation
50+
return maybeFoldList(nonNullStream.map(this::convert).toList());
51+
}
52+
53+
private Object maybeFoldList(List<?> list) {
54+
if (list.isEmpty()) {
55+
return null;
56+
}
57+
58+
if (list.size() == 1) {
59+
return list.get(0);
60+
}
61+
62+
return list;
63+
}
64+
65+
private BytesRef convert(String value) {
66+
return new BytesRef(value);
67+
}
68+
}

‎test/framework/src/main/java/org/elasticsearch/index/mapper/BlockLoaderTestCase.java‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.index.mapper;
11+
12+
import org.apache.lucene.index.DirectoryReader;
13+
import org.apache.lucene.index.LeafReaderContext;
14+
import org.apache.lucene.store.Directory;
15+
import org.apache.lucene.tests.index.RandomIndexWriter;
16+
import org.elasticsearch.index.IndexSettings;
17+
import org.elasticsearch.index.fieldvisitor.StoredFieldLoader;
18+
import org.elasticsearch.logsdb.datageneration.DataGeneratorSpecification;
19+
import org.elasticsearch.logsdb.datageneration.FieldDataGenerator;
20+
import org.elasticsearch.logsdb.datageneration.FieldType;
21+
import org.elasticsearch.logsdb.datageneration.MappingGenerator;
22+
import org.elasticsearch.logsdb.datageneration.Template;
23+
import org.elasticsearch.logsdb.datageneration.datasource.DataSourceHandler;
24+
import org.elasticsearch.logsdb.datageneration.datasource.DataSourceRequest;
25+
import org.elasticsearch.logsdb.datageneration.datasource.DataSourceResponse;
26+
import org.elasticsearch.search.fetch.StoredFieldsSpec;
27+
import org.elasticsearch.search.lookup.SearchLookup;
28+
import org.elasticsearch.xcontent.XContentBuilder;
29+
import org.elasticsearch.xcontent.XContentType;
30+
31+
import java.io.IOException;
32+
import java.util.List;
33+
import java.util.Map;
34+
import java.util.Set;
35+
36+
public abstract class BlockLoaderTestCase extends MapperServiceTestCase {
37+
private final String fieldName;
38+
private final Template template;
39+
private final MappingGenerator mappingGenerator;
40+
private final FieldDataGenerator generator;
41+
42+
protected BlockLoaderTestCase(FieldType fieldType) {
43+
this.fieldName = randomAlphaOfLengthBetween(5, 10);
44+
45+
// Disable all dynamic mapping
46+
var specification = DataGeneratorSpecification.builder()
47+
.withFullyDynamicMapping(false)
48+
.withDataSourceHandlers(List.of(new DataSourceHandler() {
49+
@Override
50+
public DataSourceResponse.DynamicMappingGenerator handle(DataSourceRequest.DynamicMappingGenerator request) {
51+
return new DataSourceResponse.DynamicMappingGenerator(isObject -> false);
52+
}
53+
}))
54+
.build();
55+
56+
this.template = new Template(Map.of(fieldName, new Template.Leaf(fieldName, fieldType)));
57+
this.mappingGenerator = new MappingGenerator(specification);
58+
this.generator = fieldType.generator(fieldName, specification.dataSource());
59+
}
60+
61+
public void testBlockLoader() throws IOException {
62+
var mapping = mappingGenerator.generate(template);
63+
var mappingXContent = XContentBuilder.builder(XContentType.JSON.xContent()).map(mapping.raw());
64+
65+
var syntheticSource = randomBoolean();
66+
var mapperService = syntheticSource ? createSytheticSourceMapperService(mappingXContent) : createMapperService(mappingXContent);
67+
68+
var fieldValue = generator.generateValue();
69+
70+
Object blockLoaderResult = setupAndInvokeBlockLoader(mapperService, fieldValue);
71+
Object expected = expected(mapping.lookup().get(fieldName), fieldValue, syntheticSource);
72+
assertEquals(expected, blockLoaderResult);
73+
}
74+
75+
protected abstract Object expected(Map<String, Object> fieldMapping, Object value, boolean syntheticSource);
76+
77+
private Object setupAndInvokeBlockLoader(MapperService mapperService, Object fieldValue) throws IOException {
78+
try (Directory directory = newDirectory()) {
79+
RandomIndexWriter iw = new RandomIndexWriter(random(), directory);
80+
81+
LuceneDocument doc = mapperService.documentMapper().parse(source(b -> {
82+
b.field(fieldName);
83+
b.value(fieldValue);
84+
})).rootDoc();
85+
86+
iw.addDocument(doc);
87+
iw.close();
88+
89+
try (DirectoryReader reader = DirectoryReader.open(directory)) {
90+
LeafReaderContext context = reader.leaves().get(0);
91+
return load(createBlockLoader(mapperService), context, mapperService);
92+
}
93+
}
94+
}
95+
96+
private Object load(BlockLoader blockLoader, LeafReaderContext context, MapperService mapperService) throws IOException {
97+
// `columnAtATimeReader` is tried first, we mimic `ValuesSourceReaderOperator`
98+
var columnAtATimeReader = blockLoader.columnAtATimeReader(context);
99+
if (columnAtATimeReader != null) {
100+
var block = (TestBlock) columnAtATimeReader.read(TestBlock.factory(context.reader().numDocs()), TestBlock.docs(0));
101+
return block.get(0);
102+
}
103+
104+
StoredFieldsSpec storedFieldsSpec = blockLoader.rowStrideStoredFieldSpec();
105+
SourceLoader.Leaf leafSourceLoader = null;
106+
if (storedFieldsSpec.requiresSource()) {
107+
var sourceLoader = mapperService.mappingLookup().newSourceLoader(null, SourceFieldMetrics.NOOP);
108+
leafSourceLoader = sourceLoader.leaf(context.reader(), null);
109+
storedFieldsSpec = storedFieldsSpec.merge(
110+
new StoredFieldsSpec(true, storedFieldsSpec.requiresMetadata(), sourceLoader.requiredStoredFields())
111+
);
112+
}
113+
BlockLoaderStoredFieldsFromLeafLoader storedFieldsLoader = new BlockLoaderStoredFieldsFromLeafLoader(
114+
StoredFieldLoader.fromSpec(storedFieldsSpec).getLoader(context, null),
115+
leafSourceLoader
116+
);
117+
storedFieldsLoader.advanceTo(0);
118+
119+
BlockLoader.Builder builder = blockLoader.builder(TestBlock.factory(context.reader().numDocs()), 1);
120+
blockLoader.rowStrideReader(context).read(0, storedFieldsLoader, builder);
121+
var block = (TestBlock) builder.build();
122+
return block.get(0);
123+
}
124+
125+
private BlockLoader createBlockLoader(MapperService mapperService) {
126+
SearchLookup searchLookup = new SearchLookup(mapperService.mappingLookup().fieldTypesLookup()::get, null, null);
127+
128+
return mapperService.fieldType(fieldName).blockLoader(new MappedFieldType.BlockLoaderContext() {
129+
@Override
130+
public String indexName() {
131+
return mapperService.getIndexSettings().getIndex().getName();
132+
}
133+
134+
@Override
135+
public IndexSettings indexSettings() {
136+
return mapperService.getIndexSettings();
137+
}
138+
139+
@Override
140+
public MappedFieldType.FieldExtractPreference fieldExtractPreference() {
141+
// TODO randomize when adding support for fields that care about this
142+
return MappedFieldType.FieldExtractPreference.NONE;
143+
}
144+
145+
@Override
146+
public SearchLookup lookup() {
147+
return searchLookup;
148+
}
149+
150+
@Override
151+
public Set<String> sourcePaths(String name) {
152+
return mapperService.mappingLookup().sourcePaths(name);
153+
}
154+
155+
@Override
156+
public String parentField(String field) {
157+
return mapperService.mappingLookup().parentField(field);
158+
}
159+
160+
@Override
161+
public FieldNamesFieldMapper.FieldNamesFieldType fieldNames() {
162+
return (FieldNamesFieldMapper.FieldNamesFieldType) mapperService.fieldType(FieldNamesFieldMapper.NAME);
163+
}
164+
});
165+
}
166+
}

‎test/framework/src/main/java/org/elasticsearch/logsdb/datageneration/Template.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,9 @@
1919
* @param template actual template data
2020
*/
2121
public record Template(Map<String, Entry> template) {
22-
sealed interface Entry permits Leaf, Object {}
22+
public sealed interface Entry permits Leaf, Object {}
2323

24-
record Leaf(String name, FieldType type) implements Entry {}
24+
public record Leaf(String name, FieldType type) implements Entry {}
2525

26-
record Object(String name, boolean nested, Map<String, Entry> children) implements Entry {}
26+
public record Object(String name, boolean nested, Map<String, Entry> children) implements Entry {}
2727
}

0 commit comments

Comments
 (0)