Skip to content

Commit 3cad0e8

Browse files
Kubik42Kubik42
committed
Improve block loader for source only runtime IP fields
1 parent cbc96f2 commit 3cad0e8

File tree

2 files changed

+189
-0
lines changed

2 files changed

+189
-0
lines changed

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

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,9 @@
3333
import org.elasticsearch.search.runtime.IpScriptFieldRangeQuery;
3434
import org.elasticsearch.search.runtime.IpScriptFieldTermQuery;
3535
import org.elasticsearch.search.runtime.IpScriptFieldTermsQuery;
36+
import org.elasticsearch.xcontent.XContentParser;
3637

38+
import java.io.IOException;
3739
import java.net.InetAddress;
3840
import java.time.ZoneId;
3941
import java.util.ArrayList;
@@ -213,6 +215,50 @@ private Query cidrQuery(String term, SearchExecutionContext context) {
213215

214216
@Override
215217
public BlockLoader blockLoader(BlockLoaderContext blContext) {
218+
FallbackSyntheticSourceBlockLoader fallbackSyntheticSourceBlockLoader = fallbackSyntheticSourceBlockLoader(
219+
blContext,
220+
BlockLoader.BlockFactory::bytesRefs,
221+
this::fallbackSyntheticSourceBlockLoaderReader
222+
);
223+
224+
if (fallbackSyntheticSourceBlockLoader != null) {
225+
return fallbackSyntheticSourceBlockLoader;
226+
}
216227
return new IpScriptBlockDocValuesReader.IpScriptBlockLoader(leafFactory(blContext.lookup()));
217228
}
229+
230+
private FallbackSyntheticSourceBlockLoader.Reader<?> fallbackSyntheticSourceBlockLoaderReader() {
231+
return new FallbackSyntheticSourceBlockLoader.SingleValueReader<InetAddress>(null) {
232+
@Override
233+
public void convertValue(Object value, List<InetAddress> accumulator) {
234+
try {
235+
if (value instanceof InetAddress ia) {
236+
accumulator.add(ia);
237+
} else {
238+
accumulator.add(InetAddresses.forString(value.toString()));
239+
}
240+
} catch (Exception e) {
241+
// value is malformed, skip it
242+
}
243+
}
244+
245+
@Override
246+
public void writeToBlock(List<InetAddress> values, BlockLoader.Builder blockBuilder) {
247+
BlockLoader.BytesRefBuilder builder = (BlockLoader.BytesRefBuilder) blockBuilder;
248+
for (InetAddress addr : values) {
249+
builder.appendBytesRef(new BytesRef(InetAddressPoint.encode(addr)));
250+
}
251+
}
252+
253+
@Override
254+
protected void parseNonNullValue(XContentParser parser, List<InetAddress> accumulator) throws IOException {
255+
try {
256+
accumulator.add(InetAddresses.forString(parser.text()));
257+
} catch (Exception e) {
258+
// value is malformed, skip it
259+
}
260+
}
261+
};
262+
}
263+
218264
}

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

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import org.apache.lucene.util.BytesRef;
3131
import org.elasticsearch.common.lucene.search.function.ScriptScoreQuery;
3232
import org.elasticsearch.common.network.InetAddresses;
33+
import org.elasticsearch.common.settings.Settings;
3334
import org.elasticsearch.index.IndexVersion;
3435
import org.elasticsearch.index.fielddata.BinaryScriptFieldData;
3536
import org.elasticsearch.index.fielddata.ScriptDocValues.Strings;
@@ -43,20 +44,27 @@
4344
import org.elasticsearch.script.ScriptType;
4445
import org.elasticsearch.search.DocValueFormat;
4546
import org.elasticsearch.search.MultiValueMode;
47+
import org.elasticsearch.search.lookup.SearchLookup;
4648

4749
import java.io.IOException;
4850
import java.time.ZoneId;
4951
import java.util.ArrayList;
52+
import java.util.Arrays;
5053
import java.util.List;
5154
import java.util.Map;
5255

5356
import static java.util.Collections.emptyMap;
5457
import static org.hamcrest.Matchers.containsInAnyOrder;
5558
import static org.hamcrest.Matchers.equalTo;
59+
import static org.hamcrest.Matchers.instanceOf;
60+
import static org.hamcrest.Matchers.nullValue;
5661
import static org.hamcrest.Matchers.sameInstance;
5762

5863
public class IpScriptFieldTypeTests extends AbstractScriptFieldTypeTestCase {
5964

65+
private static final BytesRef EMPTY_IP = null;
66+
private static final BytesRef MALFORMED_IP = null;
67+
6068
@Override
6169
protected ScriptFactory parseFromSource() {
6270
return IpFieldScript.PARSE_FROM_SOURCE;
@@ -280,6 +288,141 @@ public void testBlockLoader() throws IOException {
280288
}
281289
}
282290

291+
public void testBlockLoaderSourceOnlyRuntimeField() throws IOException {
292+
try (
293+
Directory directory = newDirectory();
294+
RandomIndexWriter iw = new RandomIndexWriter(random(), directory, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))
295+
) {
296+
// given
297+
// try multiple variations of boolean as they're all encoded slightly differently
298+
iw.addDocuments(
299+
List.of(
300+
List.of(new StoredField("_source", new BytesRef("{\"test\": [\"192.168.0.1\"]}"))),
301+
List.of(new StoredField("_source", new BytesRef("{\"test\": [\"2001:db8::1\"]}"))),
302+
List.of(new StoredField("_source", new BytesRef("{\"test\": [\"\"]}"))),
303+
// ensure a malformed value doesn't crash
304+
List.of(new StoredField("_source", new BytesRef("{\"test\": [\"potato\"]}")))
305+
)
306+
);
307+
IpScriptFieldType fieldType = simpleSourceOnlyMappedFieldType();
308+
List<BytesRef> expected = Arrays.asList(
309+
new BytesRef(InetAddressPoint.encode(InetAddresses.forString("192.168.0.1"))),
310+
new BytesRef(InetAddressPoint.encode(InetAddresses.forString("2001:db8::1"))),
311+
null,
312+
MALFORMED_IP
313+
);
314+
315+
try (DirectoryReader reader = iw.getReader()) {
316+
// when
317+
BlockLoader loader = fieldType.blockLoader(blContext(Settings.EMPTY, true));
318+
319+
// then
320+
321+
// assert loader is of expected instance type
322+
assertThat(loader, instanceOf(IpScriptBlockDocValuesReader.IpScriptBlockLoader.class));
323+
324+
// ignored source doesn't support column at a time loading:
325+
var columnAtATimeLoader = loader.columnAtATimeReader(reader.leaves().getFirst());
326+
assertThat(columnAtATimeLoader, instanceOf(IpScriptBlockDocValuesReader.class));
327+
328+
var rowStrideReader = loader.rowStrideReader(reader.leaves().getFirst());
329+
assertThat(rowStrideReader, instanceOf(IpScriptBlockDocValuesReader.class));
330+
331+
// assert values
332+
assertThat(blockLoaderReadValuesFromColumnAtATimeReader(reader, fieldType, 0), equalTo(expected));
333+
assertThat(blockLoaderReadValuesFromRowStrideReader(reader, fieldType), equalTo(expected));
334+
}
335+
}
336+
}
337+
338+
public void testBlockLoaderSourceOnlyRuntimeFieldWithSyntheticSource() throws IOException {
339+
try (
340+
Directory directory = newDirectory();
341+
RandomIndexWriter iw = new RandomIndexWriter(random(), directory, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))
342+
) {
343+
// given
344+
// try multiple variations of boolean as they're all encoded slightly differently
345+
iw.addDocuments(
346+
List.of(
347+
createDocumentWithIgnoredSource("[\"192.168.0.1\"]"),
348+
createDocumentWithIgnoredSource("[\"2001:db8::1\"]"),
349+
createDocumentWithIgnoredSource("[\"\"]"),
350+
// ensure a malformed value doesn't crash
351+
createDocumentWithIgnoredSource("[\"potato\"]")
352+
)
353+
);
354+
355+
Settings settings = Settings.builder().put("index.mapping.source.mode", "synthetic").build();
356+
IpScriptFieldType fieldType = simpleSourceOnlyMappedFieldType();
357+
List<BytesRef> expected = Arrays.asList(
358+
new BytesRef(InetAddressPoint.encode(InetAddresses.forString("192.168.0.1"))),
359+
new BytesRef(InetAddressPoint.encode(InetAddresses.forString("2001:db8::1"))),
360+
EMPTY_IP,
361+
MALFORMED_IP
362+
);
363+
364+
try (DirectoryReader reader = iw.getReader()) {
365+
// when
366+
BlockLoader loader = fieldType.blockLoader(blContext(settings, true));
367+
368+
// then
369+
370+
// assert loader is of expected instance type
371+
assertThat(loader, instanceOf(FallbackSyntheticSourceBlockLoader.class));
372+
373+
// ignored source doesn't support column at a time loading:
374+
var columnAtATimeLoader = loader.columnAtATimeReader(reader.leaves().getFirst());
375+
assertThat(columnAtATimeLoader, nullValue());
376+
377+
var rowStrideReader = loader.rowStrideReader(reader.leaves().getFirst());
378+
assertThat(
379+
rowStrideReader.getClass().getName(),
380+
equalTo("org.elasticsearch.index.mapper.FallbackSyntheticSourceBlockLoader$IgnoredSourceRowStrideReader")
381+
);
382+
383+
// assert values
384+
assertThat(blockLoaderReadValuesFromRowStrideReader(settings, reader, fieldType, true), equalTo(expected));
385+
}
386+
}
387+
}
388+
389+
/**
390+
* Returns a source only mapped field type. This is useful, since the available build() function doesn't override isParsedFromSource()
391+
*/
392+
private IpScriptFieldType simpleSourceOnlyMappedFieldType() {
393+
Script script = new Script(ScriptType.INLINE, "test", "", emptyMap());
394+
IpFieldScript.Factory factory = new IpFieldScript.Factory() {
395+
@Override
396+
public IpFieldScript.LeafFactory newFactory(
397+
String fieldName,
398+
Map<String, Object> params,
399+
SearchLookup searchLookup,
400+
OnScriptError onScriptError
401+
) {
402+
return ctx -> new IpFieldScript(fieldName, params, searchLookup, onScriptError, ctx) {
403+
@Override
404+
@SuppressWarnings("unchecked")
405+
public void execute() {
406+
Map<String, Object> source = (Map<String, Object>) this.getParams().get("_source");
407+
for (Object foo : (List<?>) source.get("test")) {
408+
try {
409+
emit(foo.toString());
410+
} catch (Exception e) {
411+
// skip
412+
}
413+
}
414+
}
415+
};
416+
}
417+
418+
@Override
419+
public boolean isParsedFromSource() {
420+
return true;
421+
}
422+
};
423+
return new IpScriptFieldType("test", factory, script, emptyMap(), OnScriptError.FAIL);
424+
}
425+
283426
@Override
284427
protected Query randomTermsQuery(MappedFieldType ft, SearchExecutionContext ctx) {
285428
return ft.termsQuery(randomList(100, () -> randomIp(randomBoolean())), ctx);

0 commit comments

Comments
 (0)