Skip to content

Commit cbc96f2

Browse files
Kubik42Kubik42
authored
Improve block loader for source only runtime boolean fields (#135318)
* Improve block loader for source only runtime boolean fields * Simplify boolean casting
1 parent 5eaea78 commit cbc96f2

File tree

2 files changed

+185
-0
lines changed

2 files changed

+185
-0
lines changed

‎server/src/main/java/org/elasticsearch/index/mapper/BooleanScriptFieldType.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,12 @@
2727
import org.elasticsearch.search.lookup.SearchLookup;
2828
import org.elasticsearch.search.runtime.BooleanScriptFieldExistsQuery;
2929
import org.elasticsearch.search.runtime.BooleanScriptFieldTermQuery;
30+
import org.elasticsearch.xcontent.XContentParser;
3031

32+
import java.io.IOException;
3133
import java.time.ZoneId;
3234
import java.util.Collection;
35+
import java.util.List;
3336
import java.util.Map;
3437
import java.util.function.Function;
3538

@@ -114,9 +117,52 @@ public DocValueFormat docValueFormat(String format, ZoneId timeZone) {
114117

115118
@Override
116119
public BlockLoader blockLoader(BlockLoaderContext blContext) {
120+
FallbackSyntheticSourceBlockLoader fallbackSyntheticSourceBlockLoader = fallbackSyntheticSourceBlockLoader(
121+
blContext,
122+
BlockLoader.BlockFactory::booleans,
123+
this::fallbackSyntheticSourceBlockLoaderReader
124+
);
125+
126+
if (fallbackSyntheticSourceBlockLoader != null) {
127+
return fallbackSyntheticSourceBlockLoader;
128+
}
117129
return new BooleanScriptBlockDocValuesReader.BooleanScriptBlockLoader(leafFactory(blContext.lookup()));
118130
}
119131

132+
private FallbackSyntheticSourceBlockLoader.Reader<?> fallbackSyntheticSourceBlockLoaderReader() {
133+
return new FallbackSyntheticSourceBlockLoader.SingleValueReader<Boolean>(null) {
134+
@Override
135+
public void convertValue(Object value, List<Boolean> accumulator) {
136+
try {
137+
if (value instanceof Boolean b) {
138+
accumulator.add(b);
139+
} else {
140+
accumulator.add(Booleans.parseBoolean(value.toString(), false));
141+
}
142+
} catch (Exception e) {
143+
// value is malformed, skip it
144+
}
145+
}
146+
147+
@Override
148+
public void writeToBlock(List<Boolean> values, BlockLoader.Builder blockBuilder) {
149+
var booleanBuilder = (BlockLoader.BooleanBuilder) blockBuilder;
150+
for (boolean value : values) {
151+
booleanBuilder.appendBoolean(value);
152+
}
153+
}
154+
155+
@Override
156+
protected void parseNonNullValue(XContentParser parser, List<Boolean> accumulator) throws IOException {
157+
try {
158+
accumulator.add(parser.booleanValue());
159+
} catch (Exception e) {
160+
// value is malformed, skip it
161+
}
162+
}
163+
};
164+
}
165+
120166
@Override
121167
public BooleanScriptFieldData.Builder fielddataBuilder(FieldDataContext fieldDataContext) {
122168
return new BooleanScriptFieldData.Builder(name(), leafFactory(fieldDataContext.lookupSupplier().get()), BooleanDocValuesField::new);

‎server/src/test/java/org/elasticsearch/index/mapper/BooleanScriptFieldTypeTests.java‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@
3131
import org.apache.lucene.util.BytesRef;
3232
import org.elasticsearch.common.bytes.BytesArray;
3333
import org.elasticsearch.common.lucene.search.function.ScriptScoreQuery;
34+
import org.elasticsearch.common.settings.Settings;
35+
import org.elasticsearch.core.Booleans;
3436
import org.elasticsearch.index.IndexVersion;
3537
import org.elasticsearch.index.fielddata.BooleanScriptFieldData;
3638
import org.elasticsearch.index.fielddata.ScriptDocValues;
@@ -44,6 +46,7 @@
4446
import org.elasticsearch.script.ScriptType;
4547
import org.elasticsearch.script.field.BooleanDocValuesField;
4648
import org.elasticsearch.search.MultiValueMode;
49+
import org.elasticsearch.search.lookup.SearchLookup;
4750
import org.elasticsearch.test.ESTestCase;
4851
import org.elasticsearch.xcontent.XContentParser;
4952
import org.elasticsearch.xcontent.XContentParser.Token;
@@ -52,6 +55,7 @@
5255

5356
import java.io.IOException;
5457
import java.util.ArrayList;
58+
import java.util.Arrays;
5559
import java.util.List;
5660
import java.util.Map;
5761

@@ -60,9 +64,13 @@
6064
import static org.hamcrest.Matchers.containsInAnyOrder;
6165
import static org.hamcrest.Matchers.equalTo;
6266
import static org.hamcrest.Matchers.instanceOf;
67+
import static org.hamcrest.Matchers.nullValue;
6368

6469
public class BooleanScriptFieldTypeTests extends AbstractNonTextScriptFieldTypeTestCase {
6570

71+
private static final Boolean MALFORMED_BOOLEAN = null;
72+
private static final Boolean EMPTY_STR_BOOLEAN = false;
73+
6674
@Override
6775
protected ScriptFactory parseFromSource() {
6876
return BooleanFieldScript.PARSE_FROM_SOURCE;
@@ -453,6 +461,137 @@ public void testBlockLoader() throws IOException {
453461
}
454462
}
455463

464+
public void testBlockLoaderSourceOnlyRuntimeField() throws IOException {
465+
try (
466+
Directory directory = newDirectory();
467+
RandomIndexWriter iw = new RandomIndexWriter(random(), directory, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))
468+
) {
469+
// given
470+
// try multiple variations of boolean as they're all encoded slightly differently
471+
iw.addDocuments(
472+
List.of(
473+
List.of(new StoredField("_source", new BytesRef("{\"test\": [false]}"))),
474+
List.of(new StoredField("_source", new BytesRef("{\"test\": [true]}"))),
475+
List.of(new StoredField("_source", new BytesRef("{\"test\": [\"false\"]}"))),
476+
List.of(new StoredField("_source", new BytesRef("{\"test\": [\"true\"]}"))),
477+
List.of(new StoredField("_source", new BytesRef("{\"test\": [\"\"]}"))),
478+
// ensure a malformed value doesn't crash
479+
List.of(new StoredField("_source", new BytesRef("{\"test\": [\"potato\"]}")))
480+
)
481+
);
482+
BooleanScriptFieldType fieldType = simpleSourceOnlyMappedFieldType();
483+
List<Boolean> expected = Arrays.asList(false, true, false, true, EMPTY_STR_BOOLEAN, MALFORMED_BOOLEAN);
484+
485+
try (DirectoryReader reader = iw.getReader()) {
486+
// when
487+
BlockLoader loader = fieldType.blockLoader(blContext(Settings.EMPTY, true));
488+
489+
// then
490+
491+
// assert loader is of expected instance type
492+
assertThat(loader, instanceOf(BooleanScriptBlockDocValuesReader.BooleanScriptBlockLoader.class));
493+
494+
// ignored source doesn't support column at a time loading:
495+
var columnAtATimeLoader = loader.columnAtATimeReader(reader.leaves().getFirst());
496+
assertThat(columnAtATimeLoader, instanceOf(BooleanScriptBlockDocValuesReader.class));
497+
498+
var rowStrideReader = loader.rowStrideReader(reader.leaves().getFirst());
499+
assertThat(rowStrideReader, instanceOf(BooleanScriptBlockDocValuesReader.class));
500+
501+
// assert values
502+
assertThat(blockLoaderReadValuesFromColumnAtATimeReader(reader, fieldType, 0), equalTo(expected));
503+
assertThat(blockLoaderReadValuesFromRowStrideReader(reader, fieldType), equalTo(expected));
504+
}
505+
}
506+
}
507+
508+
public void testBlockLoaderSourceOnlyRuntimeFieldWithSyntheticSource() throws IOException {
509+
try (
510+
Directory directory = newDirectory();
511+
RandomIndexWriter iw = new RandomIndexWriter(random(), directory, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))
512+
) {
513+
// given
514+
// try multiple variations of boolean as they're all encoded slightly differently
515+
iw.addDocuments(
516+
List.of(
517+
createDocumentWithIgnoredSource("false"),
518+
createDocumentWithIgnoredSource("true"),
519+
createDocumentWithIgnoredSource("[false]"),
520+
createDocumentWithIgnoredSource("[true]"),
521+
createDocumentWithIgnoredSource("[\"false\"]"),
522+
createDocumentWithIgnoredSource("[\"true\"]"),
523+
createDocumentWithIgnoredSource("[\"\"]"),
524+
// ensure a malformed value doesn't crash
525+
createDocumentWithIgnoredSource("[\"potato\"]")
526+
)
527+
);
528+
529+
Settings settings = Settings.builder().put("index.mapping.source.mode", "synthetic").build();
530+
BooleanScriptFieldType fieldType = simpleSourceOnlyMappedFieldType();
531+
List<Boolean> expected = Arrays.asList(false, true, false, true, false, true, EMPTY_STR_BOOLEAN, MALFORMED_BOOLEAN);
532+
533+
try (DirectoryReader reader = iw.getReader()) {
534+
// when
535+
BlockLoader loader = fieldType.blockLoader(blContext(settings, true));
536+
537+
// then
538+
539+
// assert loader is of expected instance type
540+
assertThat(loader, instanceOf(FallbackSyntheticSourceBlockLoader.class));
541+
542+
// ignored source doesn't support column at a time loading:
543+
var columnAtATimeLoader = loader.columnAtATimeReader(reader.leaves().getFirst());
544+
assertThat(columnAtATimeLoader, nullValue());
545+
546+
var rowStrideReader = loader.rowStrideReader(reader.leaves().getFirst());
547+
assertThat(
548+
rowStrideReader.getClass().getName(),
549+
equalTo("org.elasticsearch.index.mapper.FallbackSyntheticSourceBlockLoader$IgnoredSourceRowStrideReader")
550+
);
551+
552+
// assert values
553+
assertThat(blockLoaderReadValuesFromRowStrideReader(settings, reader, fieldType, true), equalTo(expected));
554+
}
555+
}
556+
}
557+
558+
/**
559+
* Returns a source only mapped field type. This is useful, since the available build() function doesn't override isParsedFromSource()
560+
*/
561+
private BooleanScriptFieldType simpleSourceOnlyMappedFieldType() {
562+
Script script = new Script(ScriptType.INLINE, "test", "", emptyMap());
563+
BooleanFieldScript.Factory factory = new BooleanFieldScript.Factory() {
564+
@Override
565+
public BooleanFieldScript.LeafFactory newFactory(
566+
String fieldName,
567+
Map<String, Object> params,
568+
SearchLookup searchLookup,
569+
OnScriptError onScriptError
570+
) {
571+
return ctx -> new BooleanFieldScript(fieldName, params, searchLookup, onScriptError, ctx) {
572+
@Override
573+
@SuppressWarnings("unchecked")
574+
public void execute() {
575+
Map<String, Object> source = (Map<String, Object>) this.getParams().get("_source");
576+
for (Object foo : (List<?>) source.get("test")) {
577+
try {
578+
emit(Booleans.parseBoolean(foo.toString(), false));
579+
} catch (Exception e) {
580+
// skip
581+
}
582+
}
583+
}
584+
};
585+
}
586+
587+
@Override
588+
public boolean isParsedFromSource() {
589+
return true;
590+
}
591+
};
592+
return new BooleanScriptFieldType("test", factory, script, emptyMap(), OnScriptError.FAIL);
593+
}
594+
456595
private void assertSameCount(IndexSearcher searcher, String source, Object queryDescription, Query scriptedQuery, Query ootbQuery)
457596
throws IOException {
458597
assertThat(

0 commit comments

Comments
 (0)