Skip to content

Commit 619cad9

Browse files
martijnvgmartijnvg
committed
Move score script context from SearchScript to its own class (#30816)
1 parent a364f18 commit 619cad9

File tree

9 files changed

+221
-42
lines changed

9 files changed

+221
-42
lines changed

‎modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,10 @@
2323
import org.apache.lucene.expressions.SimpleBindings;
2424
import org.apache.lucene.expressions.js.JavascriptCompiler;
2525
import org.apache.lucene.expressions.js.VariableContext;
26+
import org.apache.lucene.index.LeafReaderContext;
2627
import org.apache.lucene.queries.function.ValueSource;
2728
import org.apache.lucene.queries.function.valuesource.DoubleConstValueSource;
29+
import org.apache.lucene.search.Scorer;
2830
import org.apache.lucene.search.SortField;
2931
import org.elasticsearch.SpecialPermission;
3032
import org.elasticsearch.common.Nullable;
@@ -39,12 +41,14 @@
3941
import org.elasticsearch.script.ClassPermission;
4042
import org.elasticsearch.script.ExecutableScript;
4143
import org.elasticsearch.script.FilterScript;
44+
import org.elasticsearch.script.ScoreScript;
4245
import org.elasticsearch.script.ScriptContext;
4346
import org.elasticsearch.script.ScriptEngine;
4447
import org.elasticsearch.script.ScriptException;
4548
import org.elasticsearch.script.SearchScript;
4649
import org.elasticsearch.search.lookup.SearchLookup;
4750

51+
import java.io.IOException;
4852
import java.security.AccessControlContext;
4953
import java.security.AccessController;
5054
import java.security.PrivilegedAction;
@@ -111,6 +115,9 @@ protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundE
111115
} else if (context.instanceClazz.equals(FilterScript.class)) {
112116
FilterScript.Factory factory = (p, lookup) -> newFilterScript(expr, lookup, p);
113117
return context.factoryClazz.cast(factory);
118+
} else if (context.instanceClazz.equals(ScoreScript.class)) {
119+
ScoreScript.Factory factory = (p, lookup) -> newScoreScript(expr, lookup, p);
120+
return context.factoryClazz.cast(factory);
114121
}
115122
throw new IllegalArgumentException("expression engine does not know how to handle script context [" + context.name + "]");
116123
}
@@ -260,6 +267,42 @@ public void setDocument(int docid) {
260267
};
261268
};
262269
}
270+
271+
private ScoreScript.LeafFactory newScoreScript(Expression expr, SearchLookup lookup, @Nullable Map<String, Object> vars) {
272+
SearchScript.LeafFactory searchLeafFactory = newSearchScript(expr, lookup, vars);
273+
return new ScoreScript.LeafFactory() {
274+
@Override
275+
public boolean needs_score() {
276+
return searchLeafFactory.needs_score();
277+
}
278+
279+
@Override
280+
public ScoreScript newInstance(LeafReaderContext ctx) throws IOException {
281+
SearchScript script = searchLeafFactory.newInstance(ctx);
282+
return new ScoreScript(vars, lookup, ctx) {
283+
@Override
284+
public double execute() {
285+
return script.runAsDouble();
286+
}
287+
288+
@Override
289+
public void setDocument(int docid) {
290+
script.setDocument(docid);
291+
}
292+
293+
@Override
294+
public void setScorer(Scorer scorer) {
295+
script.setScorer(scorer);
296+
}
297+
298+
@Override
299+
public double get_score() {
300+
return script.getScore();
301+
}
302+
};
303+
}
304+
};
305+
}
263306

264307
/**
265308
* converts a ParseException at compile-time or link-time to a ScriptException

‎plugins/examples/script-expert-scoring/src/main/java/org/elasticsearch/example/expertscript/ExpertScriptPlugin.java

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,9 @@
3030
import org.elasticsearch.common.settings.Settings;
3131
import org.elasticsearch.plugins.Plugin;
3232
import org.elasticsearch.plugins.ScriptPlugin;
33+
import org.elasticsearch.script.ScoreScript;
3334
import org.elasticsearch.script.ScriptContext;
3435
import org.elasticsearch.script.ScriptEngine;
35-
import org.elasticsearch.script.SearchScript;
3636

3737
/**
3838
* An example script plugin that adds a {@link ScriptEngine} implementing expert scoring.
@@ -54,12 +54,12 @@ public String getType() {
5454

5555
@Override
5656
public <T> T compile(String scriptName, String scriptSource, ScriptContext<T> context, Map<String, String> params) {
57-
if (context.equals(SearchScript.SCRIPT_SCORE_CONTEXT) == false) {
57+
if (context.equals(ScoreScript.CONTEXT) == false) {
5858
throw new IllegalArgumentException(getType() + " scripts cannot be used for context [" + context.name + "]");
5959
}
6060
// we use the script "source" as the script identifier
6161
if ("pure_df".equals(scriptSource)) {
62-
SearchScript.Factory factory = (p, lookup) -> new SearchScript.LeafFactory() {
62+
ScoreScript.Factory factory = (p, lookup) -> new ScoreScript.LeafFactory() {
6363
final String field;
6464
final String term;
6565
{
@@ -74,18 +74,18 @@ public <T> T compile(String scriptName, String scriptSource, ScriptContext<T> co
7474
}
7575

7676
@Override
77-
public SearchScript newInstance(LeafReaderContext context) throws IOException {
77+
public ScoreScript newInstance(LeafReaderContext context) throws IOException {
7878
PostingsEnum postings = context.reader().postings(new Term(field, term));
7979
if (postings == null) {
8080
// the field and/or term don't exist in this segment, so always return 0
81-
return new SearchScript(p, lookup, context) {
81+
return new ScoreScript(p, lookup, context) {
8282
@Override
83-
public double runAsDouble() {
83+
public double execute() {
8484
return 0.0d;
8585
}
8686
};
8787
}
88-
return new SearchScript(p, lookup, context) {
88+
return new ScoreScript(p, lookup, context) {
8989
int currentDocid = -1;
9090
@Override
9191
public void setDocument(int docid) {
@@ -100,7 +100,7 @@ public void setDocument(int docid) {
100100
currentDocid = docid;
101101
}
102102
@Override
103-
public double runAsDouble() {
103+
public double execute() {
104104
if (postings.docID() != currentDocid) {
105105
// advance moved past the current doc, so this doc has no occurrences of the term
106106
return 0.0d;

‎server/src/main/java/org/elasticsearch/common/lucene/search/function/ScriptScoreFunction.java

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@
2424
import org.apache.lucene.search.Explanation;
2525
import org.apache.lucene.search.Scorer;
2626
import org.elasticsearch.script.ExplainableSearchScript;
27+
import org.elasticsearch.script.ScoreScript;
2728
import org.elasticsearch.script.Script;
28-
import org.elasticsearch.script.SearchScript;
2929

3030
import java.io.IOException;
3131
import java.util.Objects;
@@ -58,18 +58,18 @@ public DocIdSetIterator iterator() {
5858

5959
private final Script sScript;
6060

61-
private final SearchScript.LeafFactory script;
61+
private final ScoreScript.LeafFactory script;
6262

6363

64-
public ScriptScoreFunction(Script sScript, SearchScript.LeafFactory script) {
64+
public ScriptScoreFunction(Script sScript, ScoreScript.LeafFactory script) {
6565
super(CombineFunction.REPLACE);
6666
this.sScript = sScript;
6767
this.script = script;
6868
}
6969

7070
@Override
7171
public LeafScoreFunction getLeafScoreFunction(LeafReaderContext ctx) throws IOException {
72-
final SearchScript leafScript = script.newInstance(ctx);
72+
final ScoreScript leafScript = script.newInstance(ctx);
7373
final CannedScorer scorer = new CannedScorer();
7474
leafScript.setScorer(scorer);
7575
return new LeafScoreFunction() {
@@ -78,7 +78,7 @@ public double score(int docId, float subQueryScore) throws IOException {
7878
leafScript.setDocument(docId);
7979
scorer.docid = docId;
8080
scorer.score = subQueryScore;
81-
double result = leafScript.runAsDouble();
81+
double result = leafScript.execute();
8282
return result;
8383
}
8484

‎server/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import org.elasticsearch.common.xcontent.XContentParser;
2929
import org.elasticsearch.index.query.QueryShardContext;
3030
import org.elasticsearch.index.query.QueryShardException;
31+
import org.elasticsearch.script.ScoreScript;
3132
import org.elasticsearch.script.Script;
3233
import org.elasticsearch.script.SearchScript;
3334

@@ -92,8 +93,8 @@ protected int doHashCode() {
9293
@Override
9394
protected ScoreFunction doToFunction(QueryShardContext context) {
9495
try {
95-
SearchScript.Factory factory = context.getScriptService().compile(script, SearchScript.SCRIPT_SCORE_CONTEXT);
96-
SearchScript.LeafFactory searchScript = factory.newFactory(script.getParams(), context.lookup());
96+
ScoreScript.Factory factory = context.getScriptService().compile(script, ScoreScript.CONTEXT);
97+
ScoreScript.LeafFactory searchScript = factory.newFactory(script.getParams(), context.lookup());
9798
return new ScriptScoreFunction(script, searchScript);
9899
} catch (Exception e) {
99100
throw new QueryShardException(context, "script_score: the script could not be loaded", e);

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

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
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 org.apache.lucene.index.LeafReaderContext;
22+
import org.apache.lucene.search.Scorer;
23+
import org.elasticsearch.index.fielddata.ScriptDocValues;
24+
import org.elasticsearch.search.lookup.LeafSearchLookup;
25+
import org.elasticsearch.search.lookup.SearchLookup;
26+
27+
import java.io.IOException;
28+
import java.io.UncheckedIOException;
29+
import java.util.Map;
30+
import java.util.function.DoubleSupplier;
31+
32+
/**
33+
* A script used for adjusting the score on a per document basis.
34+
*/
35+
public abstract class ScoreScript {
36+
37+
public static final String[] PARAMETERS = new String[]{};
38+
39+
/** The generic runtime parameters for the script. */
40+
private final Map<String, Object> params;
41+
42+
/** A leaf lookup for the bound segment this script will operate on. */
43+
private final LeafSearchLookup leafLookup;
44+
45+
private DoubleSupplier scoreSupplier = () -> 0.0;
46+
47+
public ScoreScript(Map<String, Object> params, SearchLookup lookup, LeafReaderContext leafContext) {
48+
this.params = params;
49+
this.leafLookup = lookup.getLeafSearchLookup(leafContext);
50+
}
51+
52+
public abstract double execute();
53+
54+
/** Return the parameters for this script. */
55+
public Map<String, Object> getParams() {
56+
return params;
57+
}
58+
59+
/** The doc lookup for the Lucene segment this script was created for. */
60+
public final Map<String, ScriptDocValues<?>> getDoc() {
61+
return leafLookup.doc();
62+
}
63+
64+
/** Set the current document to run the script on next. */
65+
public void setDocument(int docid) {
66+
leafLookup.setDocument(docid);
67+
}
68+
69+
public void setScorer(Scorer scorer) {
70+
this.scoreSupplier = () -> {
71+
try {
72+
return scorer.score();
73+
} catch (IOException e) {
74+
throw new UncheckedIOException(e);
75+
}
76+
};
77+
}
78+
79+
public double get_score() {
80+
return scoreSupplier.getAsDouble();
81+
}
82+
83+
/** A factory to construct {@link ScoreScript} instances. */
84+
public interface LeafFactory {
85+
86+
/**
87+
* Return {@code true} if the script needs {@code _score} calculated, or {@code false} otherwise.
88+
*/
89+
boolean needs_score();
90+
91+
ScoreScript newInstance(LeafReaderContext ctx) throws IOException;
92+
}
93+
94+
/** A factory to construct stateful {@link ScoreScript} factories for a specific index. */
95+
public interface Factory {
96+
97+
ScoreScript.LeafFactory newFactory(Map<String, Object> params, SearchLookup lookup);
98+
99+
}
100+
101+
public static final ScriptContext<ScoreScript.Factory> CONTEXT = new ScriptContext<>("score", ScoreScript.Factory.class);
102+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ public class ScriptModule {
4242
CORE_CONTEXTS = Stream.of(
4343
SearchScript.CONTEXT,
4444
SearchScript.AGGS_CONTEXT,
45-
SearchScript.SCRIPT_SCORE_CONTEXT,
45+
ScoreScript.CONTEXT,
4646
SearchScript.SCRIPT_SORT_CONTEXT,
4747
SearchScript.TERMS_SET_QUERY_CONTEXT,
4848
ExecutableScript.CONTEXT,

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -162,8 +162,6 @@ public interface Factory {
162162
public static final ScriptContext<Factory> AGGS_CONTEXT = new ScriptContext<>("aggs", Factory.class);
163163
// Can return a double. (For ScriptSortType#NUMBER only, for ScriptSortType#STRING normal CONTEXT should be used)
164164
public static final ScriptContext<Factory> SCRIPT_SORT_CONTEXT = new ScriptContext<>("sort", Factory.class);
165-
// Can return a float
166-
public static final ScriptContext<Factory> SCRIPT_SCORE_CONTEXT = new ScriptContext<>("score", Factory.class);
167165
// Can return a long
168166
public static final ScriptContext<Factory> TERMS_SET_QUERY_CONTEXT = new ScriptContext<>("terms_set", Factory.class);
169167
}

‎server/src/test/java/org/elasticsearch/search/functionscore/ExplainableScriptIT.java

Lines changed: 16 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,14 @@
3030
import org.elasticsearch.plugins.Plugin;
3131
import org.elasticsearch.plugins.ScriptPlugin;
3232
import org.elasticsearch.script.ExplainableSearchScript;
33+
import org.elasticsearch.script.ScoreScript;
3334
import org.elasticsearch.script.Script;
3435
import org.elasticsearch.script.ScriptContext;
3536
import org.elasticsearch.script.ScriptEngine;
3637
import org.elasticsearch.script.ScriptType;
37-
import org.elasticsearch.script.SearchScript;
3838
import org.elasticsearch.search.SearchHit;
3939
import org.elasticsearch.search.SearchHits;
40-
import org.elasticsearch.search.lookup.LeafDocLookup;
40+
import org.elasticsearch.search.lookup.SearchLookup;
4141
import org.elasticsearch.test.ESIntegTestCase;
4242
import org.elasticsearch.test.ESIntegTestCase.ClusterScope;
4343
import org.elasticsearch.test.ESIntegTestCase.Scope;
@@ -76,45 +76,39 @@ public String getType() {
7676
@Override
7777
public <T> T compile(String scriptName, String scriptSource, ScriptContext<T> context, Map<String, String> params) {
7878
assert scriptSource.equals("explainable_script");
79-
assert context == SearchScript.SCRIPT_SCORE_CONTEXT;
80-
SearchScript.Factory factory = (p, lookup) -> new SearchScript.LeafFactory() {
81-
@Override
82-
public SearchScript newInstance(LeafReaderContext context) throws IOException {
83-
return new MyScript(lookup.doc().getLeafDocLookup(context));
84-
}
79+
assert context == ScoreScript.CONTEXT;
80+
ScoreScript.Factory factory = (params1, lookup) -> new ScoreScript.LeafFactory() {
8581
@Override
8682
public boolean needs_score() {
8783
return false;
8884
}
85+
86+
@Override
87+
public ScoreScript newInstance(LeafReaderContext ctx) throws IOException {
88+
return new MyScript(params1, lookup, ctx);
89+
}
8990
};
9091
return context.factoryClazz.cast(factory);
9192
}
9293
};
9394
}
9495
}
9596

96-
static class MyScript extends SearchScript implements ExplainableSearchScript {
97-
LeafDocLookup docLookup;
97+
static class MyScript extends ScoreScript implements ExplainableSearchScript {
9898

99-
MyScript(LeafDocLookup docLookup) {
100-
super(null, null, null);
101-
this.docLookup = docLookup;
99+
MyScript(Map<String, Object> params, SearchLookup lookup, LeafReaderContext leafContext) {
100+
super(params, lookup, leafContext);
102101
}
103-
104-
@Override
105-
public void setDocument(int doc) {
106-
docLookup.setDocument(doc);
107-
}
108-
102+
109103
@Override
110104
public Explanation explain(Explanation subQueryScore) throws IOException {
111105
Explanation scoreExp = Explanation.match(subQueryScore.getValue(), "_score: ", subQueryScore);
112-
return Explanation.match((float) (runAsDouble()), "This script returned " + runAsDouble(), scoreExp);
106+
return Explanation.match((float) (execute()), "This script returned " + execute(), scoreExp);
113107
}
114108

115109
@Override
116-
public double runAsDouble() {
117-
return ((Number) ((ScriptDocValues) docLookup.get("number_field")).getValues().get(0)).doubleValue();
110+
public double execute() {
111+
return ((Number) ((ScriptDocValues) getDoc().get("number_field")).getValues().get(0)).doubleValue();
118112
}
119113
}
120114

0 commit comments

Comments
 (0)