Skip to content

Commit 2659648

Browse files
nik9000nik9000
authored
Add term query for keyword script fields (#59372)
This adds what I think is just about the simplest possible `term` query implementation for `keyword` script fields and wires it into the field mapper that we build for them.
1 parent fa48ccd commit 2659648

File tree

11 files changed

+291
-48
lines changed

11 files changed

+291
-48
lines changed

‎x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScript.java

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,9 @@
1414
import org.elasticsearch.search.lookup.SearchLookup;
1515

1616
import java.io.IOException;
17+
import java.util.ArrayList;
1718
import java.util.List;
1819
import java.util.Map;
19-
import java.util.function.Consumer;
2020

2121
public abstract class StringScriptFieldScript extends AbstractScriptFieldScript {
2222
public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("string_script_field", Factory.class);
@@ -32,14 +32,26 @@ public interface Factory extends ScriptFactory {
3232
}
3333

3434
public interface LeafFactory {
35-
StringScriptFieldScript newInstance(LeafReaderContext ctx, Consumer<String> sync) throws IOException;
35+
StringScriptFieldScript newInstance(LeafReaderContext ctx) throws IOException;
3636
}
3737

38-
private final Consumer<String> sync;
38+
private final List<String> results = new ArrayList<>();
3939

40-
public StringScriptFieldScript(Map<String, Object> params, SearchLookup searchLookup, LeafReaderContext ctx, Consumer<String> sync) {
40+
public StringScriptFieldScript(Map<String, Object> params, SearchLookup searchLookup, LeafReaderContext ctx) {
4141
super(params, searchLookup, ctx);
42-
this.sync = sync;
42+
}
43+
44+
/**
45+
* Execute the script for the provided {@code docId}.
46+
* <p>
47+
* @return a mutable {@link List} that contains the results of the script
48+
* and will be modified the next time you call {@linkplain #resultsForDoc}.
49+
*/
50+
public final List<String> resultsForDoc(int docId) {
51+
results.clear();
52+
setDocument(docId);
53+
execute();
54+
return results;
4355
}
4456

4557
public static class Value {
@@ -50,7 +62,7 @@ public Value(StringScriptFieldScript script) {
5062
}
5163

5264
public void value(String v) {
53-
script.sync.accept(v);
65+
script.results.add(v);
5466
}
5567
}
5668
}

‎x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/fielddata/ScriptBinaryDocValues.java

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -9,30 +9,26 @@
99
import org.elasticsearch.index.fielddata.SortingBinaryDocValues;
1010
import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
1111

12-
public final class ScriptBinaryDocValues extends SortingBinaryDocValues {
12+
import java.util.List;
1313

14+
public final class ScriptBinaryDocValues extends SortingBinaryDocValues {
1415
private final StringScriptFieldScript script;
15-
private final ScriptBinaryFieldData.ScriptBinaryResult scriptBinaryResult;
1616

17-
ScriptBinaryDocValues(StringScriptFieldScript script, ScriptBinaryFieldData.ScriptBinaryResult scriptBinaryResult) {
17+
ScriptBinaryDocValues(StringScriptFieldScript script) {
1818
this.script = script;
19-
this.scriptBinaryResult = scriptBinaryResult;
2019
}
2120

2221
@Override
23-
public boolean advanceExact(int doc) {
24-
script.setDocument(doc);
25-
script.execute();
26-
27-
count = scriptBinaryResult.getResult().size();
22+
public boolean advanceExact(int docId) {
23+
List<String> results = script.resultsForDoc(docId);
24+
count = results.size();
2825
if (count == 0) {
29-
grow();
3026
return false;
3127
}
3228

29+
grow();
3330
int i = 0;
34-
for (String value : scriptBinaryResult.getResult()) {
35-
grow();
31+
for (String value : results) {
3632
values[i++].copyChars(value);
3733
}
3834
sort();

‎x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/fielddata/ScriptBinaryFieldData.java

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -102,10 +102,7 @@ public ScriptBinaryLeafFieldData load(LeafReaderContext context) {
102102

103103
@Override
104104
public ScriptBinaryLeafFieldData loadDirect(LeafReaderContext context) throws IOException {
105-
ScriptBinaryResult scriptBinaryResult = new ScriptBinaryResult();
106-
return new ScriptBinaryLeafFieldData(
107-
new ScriptBinaryDocValues(leafFactory.get().newInstance(context, scriptBinaryResult::accept), scriptBinaryResult)
108-
);
105+
return new ScriptBinaryLeafFieldData(new ScriptBinaryDocValues(leafFactory.get().newInstance(context)));
109106
}
110107

111108
@Override

‎x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/RuntimeKeywordMappedFieldType.java

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,18 +8,21 @@
88

99
import org.apache.lucene.search.Query;
1010
import org.apache.lucene.util.BytesRef;
11-
import org.elasticsearch.common.xcontent.XContentBuilder;
1211
import org.elasticsearch.common.xcontent.ToXContent.Params;
12+
import org.elasticsearch.common.lucene.BytesRefs;
13+
import org.elasticsearch.common.xcontent.XContentBuilder;
1314
import org.elasticsearch.index.fielddata.IndexFieldData;
1415
import org.elasticsearch.index.mapper.MappedFieldType;
1516
import org.elasticsearch.index.mapper.TextSearchInfo;
1617
import org.elasticsearch.index.query.QueryShardContext;
1718
import org.elasticsearch.script.Script;
1819
import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
1920
import org.elasticsearch.xpack.runtimefields.fielddata.ScriptBinaryFieldData;
21+
import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldTermQuery;
2022

2123
import java.io.IOException;
2224
import java.util.Map;
25+
import java.util.Objects;
2326

2427
public final class RuntimeKeywordMappedFieldType extends MappedFieldType {
2528

@@ -57,7 +60,11 @@ public IndexFieldData.Builder fielddataBuilder(String fullyQualifiedIndexName) {
5760

5861
@Override
5962
public Query termQuery(Object value, QueryShardContext context) {
60-
return null;
63+
return new StringScriptFieldTermQuery(
64+
scriptFactory.newFactory(script.getParams(), context.lookup()),
65+
name(),
66+
BytesRefs.toString(Objects.requireNonNull(value))
67+
);
6168
}
6269

6370
@Override

‎x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/StringScriptFieldTermQuery.java

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
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+
* you may not use this file except in compliance with the Elastic License.
5+
*/
6+
7+
package org.elasticsearch.xpack.runtimefields.query;
8+
9+
import org.apache.lucene.index.LeafReaderContext;
10+
import org.apache.lucene.index.Term;
11+
import org.apache.lucene.search.ConstantScoreScorer;
12+
import org.apache.lucene.search.ConstantScoreWeight;
13+
import org.apache.lucene.search.DocIdSetIterator;
14+
import org.apache.lucene.search.IndexSearcher;
15+
import org.apache.lucene.search.Query;
16+
import org.apache.lucene.search.QueryVisitor;
17+
import org.apache.lucene.search.ScoreMode;
18+
import org.apache.lucene.search.Scorer;
19+
import org.apache.lucene.search.TwoPhaseIterator;
20+
import org.apache.lucene.search.Weight;
21+
import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
22+
23+
import java.io.IOException;
24+
import java.util.Objects;
25+
26+
public class StringScriptFieldTermQuery extends Query {
27+
private final StringScriptFieldScript.LeafFactory leafFactory;
28+
private final String fieldName;
29+
private final String term;
30+
31+
public StringScriptFieldTermQuery(StringScriptFieldScript.LeafFactory leafFactory, String fieldName, String term) {
32+
this.leafFactory = leafFactory;
33+
this.fieldName = fieldName;
34+
this.term = term;
35+
}
36+
37+
@Override
38+
public final Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
39+
return new ConstantScoreWeight(this, boost) {
40+
@Override
41+
public boolean isCacheable(LeafReaderContext ctx) {
42+
return false; // scripts aren't really cacheable at this point
43+
}
44+
45+
@Override
46+
public Scorer scorer(LeafReaderContext ctx) throws IOException {
47+
StringScriptFieldScript script = leafFactory.newInstance(ctx);
48+
DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc());
49+
TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) {
50+
@Override
51+
public boolean matches() throws IOException {
52+
for (String result : script.resultsForDoc(approximation().docID())) {
53+
if (term.equals(result)) {
54+
return true;
55+
}
56+
}
57+
return false;
58+
}
59+
60+
@Override
61+
public float matchCost() {
62+
// TODO we don't have a good way of estimating the complexity of the script so we just go with 9000
63+
return 9000f;
64+
}
65+
};
66+
return new ConstantScoreScorer(this, score(), scoreMode, twoPhase);
67+
}
68+
};
69+
}
70+
71+
@Override
72+
public void visit(QueryVisitor visitor) {
73+
visitor.consumeTerms(this, new Term(fieldName, term));
74+
}
75+
76+
@Override
77+
public final String toString(String field) {
78+
if (fieldName.contentEquals(field)) {
79+
return term;
80+
}
81+
return fieldName + ":" + term;
82+
}
83+
84+
@Override
85+
public int hashCode() {
86+
return Objects.hash(fieldName, term);
87+
}
88+
89+
@Override
90+
public boolean equals(Object obj) {
91+
if (obj == null || getClass() != obj.getClass()) {
92+
return false;
93+
}
94+
StringScriptFieldTermQuery other = (StringScriptFieldTermQuery) obj;
95+
return fieldName.equals(other.fieldName) && term.equals(other.term);
96+
}
97+
}

‎x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/DoubleScriptFieldScriptTests.java

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,14 @@
1919
import org.elasticsearch.search.lookup.SearchLookup;
2020

2121
import java.io.IOException;
22+
import java.util.ArrayList;
2223
import java.util.List;
2324
import java.util.Map;
25+
import java.util.function.IntFunction;
2426

2527
import static org.hamcrest.Matchers.equalTo;
2628

2729
public class DoubleScriptFieldScriptTests extends ScriptFieldScriptTestCase<
28-
DoubleScriptFieldScript,
2930
DoubleScriptFieldScript.Factory,
3031
DoubleScriptFieldScript.LeafFactory,
3132
Double> {
@@ -112,11 +113,15 @@ protected DoubleScriptFieldScript.LeafFactory newLeafFactory(
112113
}
113114

114115
@Override
115-
protected DoubleScriptFieldScript newInstance(
116-
DoubleScriptFieldScript.LeafFactory leafFactory,
117-
LeafReaderContext context,
118-
List<Double> result
119-
) throws IOException {
120-
return leafFactory.newInstance(context, result::add);
116+
protected IntFunction<List<Double>> newInstance(DoubleScriptFieldScript.LeafFactory leafFactory, LeafReaderContext context)
117+
throws IOException {
118+
List<Double> results = new ArrayList<>();
119+
DoubleScriptFieldScript script = leafFactory.newInstance(context, results::add);
120+
return docId -> {
121+
results.clear();
122+
script.setDocument(docId);
123+
script.execute();
124+
return results;
125+
};
121126
}
122127
}

‎x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/LongScriptFieldScriptTests.java

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,14 @@
1818
import org.elasticsearch.search.lookup.SearchLookup;
1919

2020
import java.io.IOException;
21+
import java.util.ArrayList;
2122
import java.util.List;
2223
import java.util.Map;
24+
import java.util.function.IntFunction;
2325

2426
import static org.hamcrest.Matchers.equalTo;
2527

2628
public class LongScriptFieldScriptTests extends ScriptFieldScriptTestCase<
27-
LongScriptFieldScript,
2829
LongScriptFieldScript.Factory,
2930
LongScriptFieldScript.LeafFactory,
3031
Long> {
@@ -98,9 +99,15 @@ protected LongScriptFieldScript.LeafFactory newLeafFactory(
9899
}
99100

100101
@Override
101-
protected LongScriptFieldScript newInstance(LongScriptFieldScript.LeafFactory leafFactory, LeafReaderContext context, List<Long> result)
102+
protected IntFunction<List<Long>> newInstance(LongScriptFieldScript.LeafFactory leafFactory, LeafReaderContext context)
102103
throws IOException {
103-
104-
return leafFactory.newInstance(context, result::add);
104+
List<Long> results = new ArrayList<>();
105+
LongScriptFieldScript script = leafFactory.newInstance(context, results::add);
106+
return docId -> {
107+
results.clear();
108+
script.setDocument(docId);
109+
script.execute();
110+
return results;
111+
};
105112
}
106113
}

‎x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/ScriptFieldScriptTestCase.java

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -42,16 +42,17 @@
4242
import java.util.List;
4343
import java.util.Map;
4444
import java.util.function.Function;
45+
import java.util.function.IntFunction;
4546

4647
import static org.mockito.Mockito.mock;
4748
import static org.mockito.Mockito.when;
4849

49-
public abstract class ScriptFieldScriptTestCase<S extends AbstractScriptFieldScript, F, LF, R> extends ESTestCase {
50+
public abstract class ScriptFieldScriptTestCase<F, LF, R> extends ESTestCase {
5051
protected abstract ScriptContext<F> scriptContext();
5152

5253
protected abstract LF newLeafFactory(F factory, Map<String, Object> params, SearchLookup searchLookup);
5354

54-
protected abstract S newInstance(LF leafFactory, LeafReaderContext context, List<R> results) throws IOException;
55+
protected abstract IntFunction<List<R>> newInstance(LF leafFactory, LeafReaderContext context) throws IOException;
5556

5657
protected final List<R> execute(CheckedConsumer<RandomIndexWriter, IOException> indexBuilder, String script, MappedFieldType... types)
5758
throws IOException {
@@ -92,15 +93,14 @@ public ScoreMode scoreMode() {
9293

9394
@Override
9495
public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
95-
S compiled = newInstance(leafFactory, context, result);
96+
IntFunction<List<R>> compiled = newInstance(leafFactory, context);
9697
return new LeafCollector() {
9798
@Override
9899
public void setScorer(Scorable scorer) {}
99100

100101
@Override
101-
public void collect(int doc) {
102-
compiled.setDocument(doc);
103-
compiled.execute();
102+
public void collect(int docId) {
103+
result.addAll(compiled.apply(docId));
104104
}
105105
};
106106
}

‎x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScriptTests.java

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,11 @@
1919
import java.io.IOException;
2020
import java.util.List;
2121
import java.util.Map;
22+
import java.util.function.IntFunction;
2223

2324
import static org.hamcrest.Matchers.equalTo;
2425

2526
public class StringScriptFieldScriptTests extends ScriptFieldScriptTestCase<
26-
StringScriptFieldScript,
2727
StringScriptFieldScript.Factory,
2828
StringScriptFieldScript.LeafFactory,
2929
String> {
@@ -104,11 +104,8 @@ protected StringScriptFieldScript.LeafFactory newLeafFactory(
104104
}
105105

106106
@Override
107-
protected StringScriptFieldScript newInstance(
108-
StringScriptFieldScript.LeafFactory leafFactory,
109-
LeafReaderContext context,
110-
List<String> result
111-
) throws IOException {
112-
return leafFactory.newInstance(context, result::add);
107+
protected IntFunction<List<String>> newInstance(StringScriptFieldScript.LeafFactory leafFactory, LeafReaderContext context)
108+
throws IOException {
109+
return leafFactory.newInstance(context)::resultsForDoc;
113110
}
114111
}

0 commit comments

Comments
 (0)