Skip to content

Commit 719ef1f

Browse files
nik9000nik9000
authored
Remaining queries for script keyword fields (#59630)
Adds the remaining queries for scripted keyword fields and handles a few leftover TODOs, mostly around `hashCode` and `equals`.
1 parent c9b9b73 commit 719ef1f

22 files changed

+1269
-36
lines changed

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

+70-3
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,13 @@
66

77
package org.elasticsearch.xpack.runtimefields.mapper;
88

9+
import org.apache.lucene.search.MultiTermQuery.RewriteMethod;
910
import org.apache.lucene.search.Query;
1011
import org.apache.lucene.util.BytesRef;
12+
import org.elasticsearch.common.geo.ShapeRelation;
1113
import org.elasticsearch.common.lucene.BytesRefs;
14+
import org.elasticsearch.common.time.DateMathParser;
15+
import org.elasticsearch.common.unit.Fuzziness;
1216
import org.elasticsearch.common.xcontent.ToXContent.Params;
1317
import org.elasticsearch.common.xcontent.XContentBuilder;
1418
import org.elasticsearch.index.mapper.MappedFieldType;
@@ -18,10 +22,16 @@
1822
import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
1923
import org.elasticsearch.xpack.runtimefields.fielddata.ScriptBinaryFieldData;
2024
import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldExistsQuery;
25+
import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldFuzzyQuery;
26+
import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldPrefixQuery;
27+
import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldRangeQuery;
28+
import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldRegexpQuery;
2129
import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldTermQuery;
2230
import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldTermsQuery;
31+
import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldWildcardQuery;
2332

2433
import java.io.IOException;
34+
import java.time.ZoneId;
2535
import java.util.List;
2636
import java.util.Map;
2737
import java.util.Objects;
@@ -69,18 +79,75 @@ private StringScriptFieldScript.LeafFactory leafFactory(QueryShardContext contex
6979

7080
@Override
7181
public Query existsQuery(QueryShardContext context) {
72-
return new StringScriptFieldExistsQuery(leafFactory(context), name());
82+
return new StringScriptFieldExistsQuery(script, leafFactory(context), name());
83+
}
84+
85+
@Override
86+
public Query fuzzyQuery(
87+
Object value,
88+
Fuzziness fuzziness,
89+
int prefixLength,
90+
int maxExpansions,
91+
boolean transpositions,
92+
QueryShardContext context
93+
) {
94+
return StringScriptFieldFuzzyQuery.build(
95+
script,
96+
leafFactory(context),
97+
name(),
98+
BytesRefs.toString(Objects.requireNonNull(value)),
99+
fuzziness.asDistance(BytesRefs.toString(value)),
100+
prefixLength,
101+
transpositions
102+
);
103+
}
104+
105+
@Override
106+
public Query prefixQuery(String value, RewriteMethod method, org.elasticsearch.index.query.QueryShardContext context) {
107+
return new StringScriptFieldPrefixQuery(script, leafFactory(context), name(), value);
108+
}
109+
110+
@Override
111+
public Query rangeQuery(
112+
Object lowerTerm,
113+
Object upperTerm,
114+
boolean includeLower,
115+
boolean includeUpper,
116+
ShapeRelation relation,
117+
ZoneId timeZone,
118+
DateMathParser parser,
119+
QueryShardContext context
120+
) {
121+
return new StringScriptFieldRangeQuery(
122+
script,
123+
leafFactory(context),
124+
name(),
125+
BytesRefs.toString(Objects.requireNonNull(lowerTerm)),
126+
BytesRefs.toString(Objects.requireNonNull(upperTerm)),
127+
includeLower,
128+
includeUpper
129+
);
130+
}
131+
132+
@Override
133+
public Query regexpQuery(String value, int flags, int maxDeterminizedStates, RewriteMethod method, QueryShardContext context) {
134+
return new StringScriptFieldRegexpQuery(script, leafFactory(context), name(), value, flags, maxDeterminizedStates);
73135
}
74136

75137
@Override
76138
public Query termQuery(Object value, QueryShardContext context) {
77-
return new StringScriptFieldTermQuery(leafFactory(context), name(), BytesRefs.toString(Objects.requireNonNull(value)));
139+
return new StringScriptFieldTermQuery(script, leafFactory(context), name(), BytesRefs.toString(Objects.requireNonNull(value)));
78140
}
79141

80142
@Override
81143
public Query termsQuery(List<?> values, QueryShardContext context) {
82144
Set<String> terms = values.stream().map(v -> BytesRefs.toString(Objects.requireNonNull(v))).collect(toSet());
83-
return new StringScriptFieldTermsQuery(leafFactory(context), name(), terms);
145+
return new StringScriptFieldTermsQuery(script, leafFactory(context), name(), terms);
146+
}
147+
148+
@Override
149+
public Query wildcardQuery(String value, RewriteMethod method, QueryShardContext context) {
150+
return new StringScriptFieldWildcardQuery(script, leafFactory(context), name(), value);
84151
}
85152

86153
void doXContentBody(XContentBuilder builder, boolean includeDefaults, Params params) throws IOException {

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

+48
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
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.search.QueryVisitor;
10+
import org.apache.lucene.util.BytesRefBuilder;
11+
import org.apache.lucene.util.automaton.ByteRunAutomaton;
12+
import org.elasticsearch.script.Script;
13+
import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
14+
15+
import java.util.List;
16+
17+
public abstract class AbstractStringScriptFieldAutomatonQuery extends AbstractStringScriptFieldQuery {
18+
private final BytesRefBuilder scratch = new BytesRefBuilder();
19+
private final ByteRunAutomaton automaton;
20+
21+
public AbstractStringScriptFieldAutomatonQuery(
22+
Script script,
23+
StringScriptFieldScript.LeafFactory leafFactory,
24+
String fieldName,
25+
ByteRunAutomaton automaton
26+
) {
27+
super(script, leafFactory, fieldName);
28+
this.automaton = automaton;
29+
}
30+
31+
@Override
32+
protected final boolean matches(List<String> values) {
33+
for (String value : values) {
34+
scratch.copyChars(value);
35+
if (automaton.run(scratch.bytes(), 0, scratch.length())) {
36+
return true;
37+
}
38+
}
39+
return false;
40+
}
41+
42+
@Override
43+
public final void visit(QueryVisitor visitor) {
44+
if (visitor.acceptField(fieldName())) {
45+
visitor.consumeTermsMatching(this, fieldName(), () -> automaton);
46+
}
47+
}
48+
}

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

+11-6
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import org.apache.lucene.search.Scorer;
1717
import org.apache.lucene.search.TwoPhaseIterator;
1818
import org.apache.lucene.search.Weight;
19+
import org.elasticsearch.script.Script;
1920
import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
2021
import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript.LeafFactory;
2122

@@ -27,18 +28,20 @@
2728
* Abstract base class for building queries based on {@link StringScriptFieldScript}.
2829
*/
2930
abstract class AbstractStringScriptFieldQuery extends Query {
31+
private final Script script;
3032
private final StringScriptFieldScript.LeafFactory leafFactory;
3133
private final String fieldName;
3234

33-
AbstractStringScriptFieldQuery(LeafFactory leafFactory, String fieldName) {
35+
AbstractStringScriptFieldQuery(Script script, LeafFactory leafFactory, String fieldName) {
36+
this.script = script;
3437
this.leafFactory = Objects.requireNonNull(leafFactory);
3538
this.fieldName = Objects.requireNonNull(fieldName);
3639
}
3740

3841
/**
3942
* Does the value match this query?
4043
*/
41-
public abstract boolean matches(List<String> values);
44+
protected abstract boolean matches(List<String> values);
4245

4346
@Override
4447
public final Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
@@ -69,23 +72,25 @@ public float matchCost() {
6972
};
7073
}
7174

75+
final Script script() {
76+
return script;
77+
}
78+
7279
protected final String fieldName() {
7380
return fieldName;
7481
}
7582

7683
@Override
7784
public int hashCode() {
78-
// TODO should leafFactory be here? Something about the script probably should be!
79-
return Objects.hash(getClass(), fieldName);
85+
return Objects.hash(getClass(), script, fieldName);
8086
}
8187

8288
@Override
8389
public boolean equals(Object obj) {
8490
if (obj == null || getClass() != obj.getClass()) {
8591
return false;
8692
}
87-
// TODO should leafFactory be here? Something about the script probably should be!
8893
AbstractStringScriptFieldQuery other = (AbstractStringScriptFieldQuery) obj;
89-
return fieldName.equals(other.fieldName);
94+
return script.equals(other.script) && fieldName.equals(other.fieldName);
9095
}
9196
}

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

+6-5
Original file line numberDiff line numberDiff line change
@@ -6,26 +6,27 @@
66

77
package org.elasticsearch.xpack.runtimefields.query;
88

9+
import org.elasticsearch.script.Script;
910
import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
1011

1112
import java.util.List;
1213

1314
public class StringScriptFieldExistsQuery extends AbstractStringScriptFieldQuery {
14-
public StringScriptFieldExistsQuery(StringScriptFieldScript.LeafFactory leafFactory, String fieldName) {
15-
super(leafFactory, fieldName);
15+
public StringScriptFieldExistsQuery(Script script, StringScriptFieldScript.LeafFactory leafFactory, String fieldName) {
16+
super(script, leafFactory, fieldName);
1617
}
1718

1819
@Override
19-
public boolean matches(List<String> values) {
20+
protected boolean matches(List<String> values) {
2021
return false == values.isEmpty();
2122
}
2223

2324
@Override
2425
public final String toString(String field) {
2526
if (fieldName().contentEquals(field)) {
26-
return "*";
27+
return "ScriptFieldExists";
2728
}
28-
return fieldName() + ":*";
29+
return fieldName() + ":ScriptFieldExists";
2930
}
3031

3132
// Superclass's equals and hashCode are great for this class

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

+68
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+
* 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.Term;
10+
import org.apache.lucene.search.FuzzyQuery;
11+
import org.apache.lucene.util.automaton.ByteRunAutomaton;
12+
import org.elasticsearch.script.Script;
13+
import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
14+
15+
import java.util.Objects;
16+
17+
public class StringScriptFieldFuzzyQuery extends AbstractStringScriptFieldAutomatonQuery {
18+
public static StringScriptFieldFuzzyQuery build(
19+
Script script,
20+
StringScriptFieldScript.LeafFactory leafFactory,
21+
String fieldName,
22+
String term,
23+
int maxEdits,
24+
int prefixLength,
25+
boolean transpositions
26+
) {
27+
int maxExpansions = 1; // We don't actually expand anything so the value here doesn't matter
28+
FuzzyQuery delegate = new FuzzyQuery(new Term(fieldName, term), maxEdits, prefixLength, maxExpansions, transpositions);
29+
ByteRunAutomaton automaton = delegate.getAutomata().runAutomaton;
30+
return new StringScriptFieldFuzzyQuery(script, leafFactory, fieldName, automaton, delegate);
31+
}
32+
33+
private final FuzzyQuery delegate;
34+
35+
private StringScriptFieldFuzzyQuery(
36+
Script script,
37+
StringScriptFieldScript.LeafFactory leafFactory,
38+
String fieldName,
39+
ByteRunAutomaton automaton,
40+
FuzzyQuery delegate
41+
) {
42+
super(script, leafFactory, fieldName, automaton);
43+
this.delegate = delegate;
44+
}
45+
46+
@Override
47+
public final String toString(String field) {
48+
return delegate.toString(field);
49+
}
50+
51+
@Override
52+
public int hashCode() {
53+
return Objects.hash(super.hashCode(), delegate);
54+
}
55+
56+
@Override
57+
public boolean equals(Object obj) {
58+
if (false == super.equals(obj)) {
59+
return false;
60+
}
61+
StringScriptFieldFuzzyQuery other = (StringScriptFieldFuzzyQuery) obj;
62+
return delegate.equals(other.delegate);
63+
}
64+
65+
FuzzyQuery delegate() {
66+
return delegate;
67+
}
68+
}

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

+69
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
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.search.PrefixQuery;
10+
import org.apache.lucene.search.QueryVisitor;
11+
import org.apache.lucene.util.BytesRef;
12+
import org.apache.lucene.util.automaton.ByteRunAutomaton;
13+
import org.elasticsearch.script.Script;
14+
import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
15+
16+
import java.util.List;
17+
import java.util.Objects;
18+
19+
public class StringScriptFieldPrefixQuery extends AbstractStringScriptFieldQuery {
20+
private final String prefix;
21+
22+
public StringScriptFieldPrefixQuery(Script script, StringScriptFieldScript.LeafFactory leafFactory, String fieldName, String prefix) {
23+
super(script, leafFactory, fieldName);
24+
this.prefix = Objects.requireNonNull(prefix);
25+
}
26+
27+
@Override
28+
protected boolean matches(List<String> values) {
29+
for (String value : values) {
30+
if (value != null && value.startsWith(prefix)) {
31+
return true;
32+
}
33+
}
34+
return false;
35+
}
36+
37+
@Override
38+
public void visit(QueryVisitor visitor) {
39+
if (visitor.acceptField(fieldName())) {
40+
visitor.consumeTermsMatching(this, fieldName(), () -> new ByteRunAutomaton(PrefixQuery.toAutomaton(new BytesRef(prefix))));
41+
}
42+
}
43+
44+
@Override
45+
public final String toString(String field) {
46+
if (fieldName().contentEquals(field)) {
47+
return prefix + "*";
48+
}
49+
return fieldName() + ":" + prefix + "*";
50+
}
51+
52+
@Override
53+
public int hashCode() {
54+
return Objects.hash(super.hashCode(), prefix);
55+
}
56+
57+
@Override
58+
public boolean equals(Object obj) {
59+
if (false == super.equals(obj)) {
60+
return false;
61+
}
62+
StringScriptFieldPrefixQuery other = (StringScriptFieldPrefixQuery) obj;
63+
return prefix.equals(other.prefix);
64+
}
65+
66+
String prefix() {
67+
return prefix;
68+
}
69+
}

0 commit comments

Comments
 (0)