Skip to content

Commit 0b4960f

Browse files
original-brownbearoriginal-brownbear
authored
SCRIPTING: Move terms_set Context to its Own Class (elastic#33602)
* SCRIPTING: Move terms_set Context to its Own Class * Extracted TermsSetQueryScript * Kept mechanics close to what they were with SearchScript
1 parent 189aace commit 0b4960f

File tree

6 files changed

+240
-14
lines changed

6 files changed

+240
-14
lines changed

‎server/src/main/java/org/elasticsearch/index/query/TermsSetQueryBuilder.java

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -40,14 +40,14 @@
4040
import org.elasticsearch.index.fielddata.IndexNumericFieldData;
4141
import org.elasticsearch.index.mapper.MappedFieldType;
4242
import org.elasticsearch.script.Script;
43-
import org.elasticsearch.script.SearchScript;
4443

4544
import java.io.IOException;
4645
import java.util.ArrayList;
4746
import java.util.HashMap;
4847
import java.util.List;
4948
import java.util.Map;
5049
import java.util.Objects;
50+
import org.elasticsearch.script.TermsSetQueryScript;
5151

5252
public final class TermsSetQueryBuilder extends AbstractQueryBuilder<TermsSetQueryBuilder> {
5353

@@ -262,13 +262,12 @@ private LongValuesSource createValuesSource(QueryShardContext context) {
262262
IndexNumericFieldData fieldData = context.getForField(msmFieldType);
263263
longValuesSource = new FieldValuesSource(fieldData);
264264
} else if (minimumShouldMatchScript != null) {
265-
SearchScript.Factory factory = context.getScriptService().compile(minimumShouldMatchScript,
266-
SearchScript.TERMS_SET_QUERY_CONTEXT);
265+
TermsSetQueryScript.Factory factory = context.getScriptService().compile(minimumShouldMatchScript,
266+
TermsSetQueryScript.CONTEXT);
267267
Map<String, Object> params = new HashMap<>();
268268
params.putAll(minimumShouldMatchScript.getParams());
269269
params.put("num_terms", values.size());
270-
SearchScript.LeafFactory leafFactory = factory.newFactory(params, context.lookup());
271-
longValuesSource = new ScriptLongValueSource(minimumShouldMatchScript, leafFactory);
270+
longValuesSource = new ScriptLongValueSource(minimumShouldMatchScript, factory.newFactory(params, context.lookup()));
272271
} else {
273272
throw new IllegalStateException("No minimum should match has been specified");
274273
}
@@ -278,26 +277,26 @@ private LongValuesSource createValuesSource(QueryShardContext context) {
278277
static final class ScriptLongValueSource extends LongValuesSource {
279278

280279
private final Script script;
281-
private final SearchScript.LeafFactory leafFactory;
280+
private final TermsSetQueryScript.LeafFactory leafFactory;
282281

283-
ScriptLongValueSource(Script script, SearchScript.LeafFactory leafFactory) {
282+
ScriptLongValueSource(Script script, TermsSetQueryScript.LeafFactory leafFactory) {
284283
this.script = script;
285284
this.leafFactory = leafFactory;
286285
}
287286

288287
@Override
289288
public LongValues getValues(LeafReaderContext ctx, DoubleValues scores) throws IOException {
290-
SearchScript searchScript = leafFactory.newInstance(ctx);
289+
TermsSetQueryScript script = leafFactory.newInstance(ctx);
291290
return new LongValues() {
292291
@Override
293292
public long longValue() throws IOException {
294-
return searchScript.runAsLong();
293+
return script.runAsLong();
295294
}
296295

297296
@Override
298297
public boolean advanceExact(int doc) throws IOException {
299-
searchScript.setDocument(doc);
300-
return searchScript.run() != null;
298+
script.setDocument(doc);
299+
return script.execute() != null;
301300
}
302301
};
303302
}

‎server/src/main/java/org/elasticsearch/script/ParameterMap.java

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.elasticsearch.script;
21+
22+
import java.util.Collection;
23+
import java.util.Map;
24+
import java.util.Set;
25+
import org.apache.logging.log4j.LogManager;
26+
import org.elasticsearch.common.logging.DeprecationLogger;
27+
28+
public final class ParameterMap implements Map<String, Object> {
29+
30+
private static final DeprecationLogger DEPRECATION_LOGGER =
31+
new DeprecationLogger(LogManager.getLogger(ParameterMap.class));
32+
33+
private final Map<String, Object> params;
34+
35+
private final Map<String, String> deprecations;
36+
37+
ParameterMap(Map<String, Object> params, Map<String, String> deprecations) {
38+
this.params = params;
39+
this.deprecations = deprecations;
40+
}
41+
42+
@Override
43+
public int size() {
44+
return params.size();
45+
}
46+
47+
@Override
48+
public boolean isEmpty() {
49+
return params.isEmpty();
50+
}
51+
52+
@Override
53+
public boolean containsKey(final Object key) {
54+
return params.containsKey(key);
55+
}
56+
57+
@Override
58+
public boolean containsValue(final Object value) {
59+
return params.containsValue(value);
60+
}
61+
62+
@Override
63+
public Object get(final Object key) {
64+
String deprecationMessage = deprecations.get(key);
65+
if (deprecationMessage != null) {
66+
DEPRECATION_LOGGER.deprecated(deprecationMessage);
67+
}
68+
return params.get(key);
69+
}
70+
71+
@Override
72+
public Object put(final String key, final Object value) {
73+
return params.put(key, value);
74+
}
75+
76+
@Override
77+
public Object remove(final Object key) {
78+
return params.remove(key);
79+
}
80+
81+
@Override
82+
public void putAll(final Map<? extends String, ?> m) {
83+
params.putAll(m);
84+
}
85+
86+
@Override
87+
public void clear() {
88+
params.clear();
89+
}
90+
91+
@Override
92+
public Set<String> keySet() {
93+
return params.keySet();
94+
}
95+
96+
@Override
97+
public Collection<Object> values() {
98+
return params.values();
99+
}
100+
101+
@Override
102+
public Set<Entry<String, Object>> entrySet() {
103+
return params.entrySet();
104+
}
105+
}

‎server/src/main/java/org/elasticsearch/script/ScriptModule.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ public class ScriptModule {
4444
SearchScript.AGGS_CONTEXT,
4545
ScoreScript.CONTEXT,
4646
SearchScript.SCRIPT_SORT_CONTEXT,
47-
SearchScript.TERMS_SET_QUERY_CONTEXT,
47+
TermsSetQueryScript.CONTEXT,
4848
ExecutableScript.CONTEXT,
4949
UpdateScript.CONTEXT,
5050
BucketAggregationScript.CONTEXT,

‎server/src/main/java/org/elasticsearch/script/SearchScript.java

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,4 @@ public interface Factory {
149149
public static final ScriptContext<Factory> AGGS_CONTEXT = new ScriptContext<>("aggs", Factory.class);
150150
// Can return a double. (For ScriptSortType#NUMBER only, for ScriptSortType#STRING normal CONTEXT should be used)
151151
public static final ScriptContext<Factory> SCRIPT_SORT_CONTEXT = new ScriptContext<>("sort", Factory.class);
152-
// Can return a long
153-
public static final ScriptContext<Factory> TERMS_SET_QUERY_CONTEXT = new ScriptContext<>("terms_set", Factory.class);
154152
}

‎server/src/main/java/org/elasticsearch/script/TermsSetQueryScript.java

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.elasticsearch.script;
20+
21+
import java.io.IOException;
22+
import java.util.Collections;
23+
import java.util.HashMap;
24+
import java.util.Map;
25+
import org.apache.lucene.index.LeafReaderContext;
26+
import org.elasticsearch.index.fielddata.ScriptDocValues;
27+
import org.elasticsearch.search.lookup.LeafSearchLookup;
28+
import org.elasticsearch.search.lookup.SearchLookup;
29+
30+
public abstract class TermsSetQueryScript {
31+
32+
public static final String[] PARAMETERS = {};
33+
34+
public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("terms_set", Factory.class);
35+
36+
private static final Map<String, String> DEPRECATIONS;
37+
38+
static {
39+
Map<String, String> deprecations = new HashMap<>();
40+
deprecations.put(
41+
"doc",
42+
"Accessing variable [doc] via [params.doc] from within a terms-set-query-script " +
43+
"is deprecated in favor of directly accessing [doc]."
44+
);
45+
deprecations.put(
46+
"_doc",
47+
"Accessing variable [doc] via [params._doc] from within a terms-set-query-script " +
48+
"is deprecated in favor of directly accessing [doc]."
49+
);
50+
DEPRECATIONS = Collections.unmodifiableMap(deprecations);
51+
}
52+
53+
/**
54+
* The generic runtime parameters for the script.
55+
*/
56+
private final Map<String, Object> params;
57+
58+
/**
59+
* A leaf lookup for the bound segment this script will operate on.
60+
*/
61+
private final LeafSearchLookup leafLookup;
62+
63+
public TermsSetQueryScript(Map<String, Object> params, SearchLookup lookup, LeafReaderContext leafContext) {
64+
this.params = new ParameterMap(params, DEPRECATIONS);
65+
this.leafLookup = lookup.getLeafSearchLookup(leafContext);
66+
}
67+
68+
/**
69+
* Return the parameters for this script.
70+
*/
71+
public Map<String, Object> getParams() {
72+
this.params.putAll(leafLookup.asMap());
73+
return params;
74+
}
75+
76+
/**
77+
* The doc lookup for the Lucene segment this script was created for.
78+
*/
79+
public Map<String, ScriptDocValues<?>> getDoc() {
80+
return leafLookup.doc();
81+
}
82+
83+
/**
84+
* Set the current document to run the script on next.
85+
*/
86+
public void setDocument(int docid) {
87+
leafLookup.setDocument(docid);
88+
}
89+
90+
/**
91+
* Return the result as a long. This is used by aggregation scripts over long fields.
92+
*/
93+
public long runAsLong() {
94+
return execute().longValue();
95+
}
96+
97+
public abstract Number execute();
98+
99+
/**
100+
* A factory to construct {@link TermsSetQueryScript} instances.
101+
*/
102+
public interface LeafFactory {
103+
TermsSetQueryScript newInstance(LeafReaderContext ctx) throws IOException;
104+
}
105+
106+
/**
107+
* A factory to construct stateful {@link TermsSetQueryScript} factories for a specific index.
108+
*/
109+
public interface Factory {
110+
LeafFactory newFactory(Map<String, Object> params, SearchLookup lookup);
111+
}
112+
}

‎test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,18 @@ public <T> T compile(String name, String source, ScriptContext<T> context, Map<S
8585
if (context.instanceClazz.equals(SearchScript.class)) {
8686
SearchScript.Factory factory = mockCompiled::createSearchScript;
8787
return context.factoryClazz.cast(factory);
88+
} else if(context.instanceClazz.equals(TermsSetQueryScript.class)) {
89+
TermsSetQueryScript.Factory factory = (parameters, lookup) -> (TermsSetQueryScript.LeafFactory) ctx
90+
-> new TermsSetQueryScript(parameters, lookup, ctx) {
91+
@Override
92+
public Number execute() {
93+
Map<String, Object> vars = new HashMap<>(parameters);
94+
vars.put("params", parameters);
95+
vars.put("doc", getDoc());
96+
return (Number) script.apply(vars);
97+
}
98+
};
99+
return context.factoryClazz.cast(factory);
88100
} else if (context.instanceClazz.equals(ExecutableScript.class)) {
89101
ExecutableScript.Factory factory = mockCompiled::createExecutableScript;
90102
return context.factoryClazz.cast(factory);

0 commit comments

Comments
 (0)