Skip to content

Commit d854edf

Browse files
original-brownbearkcm
original-brownbear
authored and
kcm
committed
SCRIPTING: Move Aggregation Script Context to its own class (#33820)
* SCRIPTING: Move Aggregation Script Context to its own class
1 parent 025550a commit d854edf

File tree

18 files changed

+384
-76
lines changed

18 files changed

+384
-76
lines changed

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

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
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.expression;
21+
22+
import java.io.IOException;
23+
import org.apache.lucene.expressions.Bindings;
24+
import org.apache.lucene.expressions.Expression;
25+
import org.apache.lucene.expressions.SimpleBindings;
26+
import org.apache.lucene.index.LeafReaderContext;
27+
import org.apache.lucene.search.DoubleValues;
28+
import org.apache.lucene.search.DoubleValuesSource;
29+
import org.elasticsearch.script.AggregationScript;
30+
import org.elasticsearch.script.GeneralScriptException;
31+
32+
/**
33+
* A bridge to evaluate an {@link Expression} against {@link Bindings} in the context
34+
* of a {@link AggregationScript}.
35+
*/
36+
class ExpressionAggregationScript implements AggregationScript.LeafFactory {
37+
38+
final Expression exprScript;
39+
final SimpleBindings bindings;
40+
final DoubleValuesSource source;
41+
final ReplaceableConstDoubleValueSource specialValue; // _value
42+
43+
ExpressionAggregationScript(Expression e, SimpleBindings b, ReplaceableConstDoubleValueSource v) {
44+
exprScript = e;
45+
bindings = b;
46+
source = exprScript.getDoubleValuesSource(bindings);
47+
specialValue = v;
48+
}
49+
50+
@Override
51+
public AggregationScript newInstance(final LeafReaderContext leaf) throws IOException {
52+
return new AggregationScript() {
53+
// Fake the scorer until setScorer is called.
54+
DoubleValues values = source.getValues(leaf, null);
55+
56+
@Override
57+
public Object execute() {
58+
try {
59+
return values.doubleValue();
60+
} catch (Exception exception) {
61+
throw new GeneralScriptException("Error evaluating " + exprScript, exception);
62+
}
63+
}
64+
65+
@Override
66+
public void setDocument(int d) {
67+
try {
68+
values.advanceExact(d);
69+
} catch (IOException e) {
70+
throw new IllegalStateException("Can't advance to doc using " + exprScript, e);
71+
}
72+
}
73+
74+
@Override
75+
public void setNextAggregationValue(Object value) {
76+
// _value isn't used in script if specialValue == null
77+
if (specialValue != null) {
78+
if (value instanceof Number) {
79+
specialValue.setValue(((Number)value).doubleValue());
80+
} else {
81+
throw new GeneralScriptException("Cannot use expression with text variable using " + exprScript);
82+
}
83+
}
84+
}
85+
};
86+
}
87+
88+
@Override
89+
public boolean needs_score() {
90+
return false;
91+
}
92+
93+
}

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import org.elasticsearch.index.mapper.DateFieldMapper;
3838
import org.elasticsearch.index.mapper.GeoPointFieldMapper.GeoPointFieldType;
3939
import org.elasticsearch.index.mapper.MappedFieldType;
40+
import org.elasticsearch.script.AggregationScript;
4041
import org.elasticsearch.script.BucketAggregationScript;
4142
import org.elasticsearch.script.BucketAggregationSelectorScript;
4243
import org.elasticsearch.script.ClassPermission;
@@ -131,6 +132,9 @@ public boolean execute() {
131132
} else if (context.instanceClazz.equals(TermsSetQueryScript.class)) {
132133
TermsSetQueryScript.Factory factory = (p, lookup) -> newTermsSetQueryScript(expr, lookup, p);
133134
return context.factoryClazz.cast(factory);
135+
} else if (context.instanceClazz.equals(AggregationScript.class)) {
136+
AggregationScript.Factory factory = (p, lookup) -> newAggregationScript(expr, lookup, p);
137+
return context.factoryClazz.cast(factory);
134138
}
135139
throw new IllegalArgumentException("expression engine does not know how to handle script context [" + context.name + "]");
136140
}
@@ -224,6 +228,37 @@ private TermsSetQueryScript.LeafFactory newTermsSetQueryScript(Expression expr,
224228
return new ExpressionTermSetQueryScript(expr, bindings);
225229
}
226230

231+
private AggregationScript.LeafFactory newAggregationScript(Expression expr, SearchLookup lookup,
232+
@Nullable Map<String, Object> vars) {
233+
// NOTE: if we need to do anything complicated with bindings in the future, we can just extend Bindings,
234+
// instead of complicating SimpleBindings (which should stay simple)
235+
SimpleBindings bindings = new SimpleBindings();
236+
ReplaceableConstDoubleValueSource specialValue = null;
237+
for (String variable : expr.variables) {
238+
try {
239+
if (variable.equals("_value")) {
240+
specialValue = new ReplaceableConstDoubleValueSource();
241+
bindings.add("_value", specialValue);
242+
// noop: _value is special for aggregations, and is handled in ExpressionScriptBindings
243+
// TODO: if some uses it in a scoring expression, they will get a nasty failure when evaluating...need a
244+
// way to know this is for aggregations and so _value is ok to have...
245+
246+
} else if (vars != null && vars.containsKey(variable)) {
247+
bindFromParams(vars, bindings, variable);
248+
} else {
249+
// delegate valuesource creation based on field's type
250+
// there are three types of "fields" to expressions, and each one has a different "api" of variables and methods.
251+
final ValueSource valueSource = getDocValueSource(variable, lookup);
252+
bindings.add(variable, valueSource.asDoubleValuesSource());
253+
}
254+
} catch (Exception e) {
255+
// we defer "binding" of variables until here: give context for that variable
256+
throw convertToScriptException("link error", expr.sourceText, variable, e);
257+
}
258+
}
259+
return new ExpressionAggregationScript(expr, bindings, specialValue);
260+
}
261+
227262
/**
228263
* This is a hack for filter scripts, which must return booleans instead of doubles as expression do.
229264
* See https://github.com/elastic/elasticsearch/issues/26429.

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

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -75,9 +75,6 @@ public boolean advanceExact(int doc) throws IOException {
7575
@Override
7676
public Object run() { return Double.valueOf(runAsDouble()); }
7777

78-
@Override
79-
public long runAsLong() { return (long)runAsDouble(); }
80-
8178
@Override
8279
public double runAsDouble() {
8380
try {

‎modules/lang-painless/src/main/java/org/elasticsearch/painless/ScriptImpl.java‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,4 @@ public double runAsDouble() {
109109
return ((Number)run()).doubleValue();
110110
}
111111

112-
@Override
113-
public long runAsLong() {
114-
return ((Number)run()).longValue();
115-
}
116112
}

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

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
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.apache.lucene.search.Scorable;
27+
import org.elasticsearch.ElasticsearchException;
28+
import org.elasticsearch.common.lucene.ScorerAware;
29+
import org.elasticsearch.index.fielddata.ScriptDocValues;
30+
import org.elasticsearch.search.lookup.LeafSearchLookup;
31+
import org.elasticsearch.search.lookup.SearchLookup;
32+
33+
public abstract class AggregationScript implements ScorerAware {
34+
35+
public static final String[] PARAMETERS = {};
36+
37+
public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("aggs", Factory.class);
38+
39+
private static final Map<String, String> DEPRECATIONS;
40+
41+
static {
42+
Map<String, String> deprecations = new HashMap<>();
43+
deprecations.put(
44+
"doc",
45+
"Accessing variable [doc] via [params.doc] from within an aggregation-script " +
46+
"is deprecated in favor of directly accessing [doc]."
47+
);
48+
deprecations.put(
49+
"_doc",
50+
"Accessing variable [doc] via [params._doc] from within an aggregation-script " +
51+
"is deprecated in favor of directly accessing [doc]."
52+
);
53+
DEPRECATIONS = Collections.unmodifiableMap(deprecations);
54+
}
55+
56+
/**
57+
* The generic runtime parameters for the script.
58+
*/
59+
private final Map<String, Object> params;
60+
61+
/**
62+
* A leaf lookup for the bound segment this script will operate on.
63+
*/
64+
private final LeafSearchLookup leafLookup;
65+
66+
/**
67+
* A scorer that will return the score for the current document when the script is run.
68+
*/
69+
protected Scorable scorer;
70+
71+
private Object value;
72+
73+
public AggregationScript(Map<String, Object> params, SearchLookup lookup, LeafReaderContext leafContext) {
74+
this.params = new ParameterMap(new HashMap<>(params), DEPRECATIONS);
75+
this.leafLookup = lookup.getLeafSearchLookup(leafContext);
76+
this.params.putAll(leafLookup.asMap());
77+
}
78+
79+
protected AggregationScript() {
80+
params = null;
81+
leafLookup = null;
82+
}
83+
84+
/**
85+
* Return the parameters for this script.
86+
*/
87+
public Map<String, Object> getParams() {
88+
return params;
89+
}
90+
91+
/**
92+
* The doc lookup for the Lucene segment this script was created for.
93+
*/
94+
public Map<String, ScriptDocValues<?>> getDoc() {
95+
return leafLookup.doc();
96+
}
97+
98+
/**
99+
* Set the current document to run the script on next.
100+
*/
101+
public void setDocument(int docid) {
102+
leafLookup.setDocument(docid);
103+
}
104+
105+
@Override
106+
public void setScorer(Scorable scorer) {
107+
this.scorer = scorer;
108+
}
109+
110+
/**
111+
* Sets per-document aggregation {@code _value}.
112+
* <p>
113+
* The default implementation just calls {@code setNextVar("_value", value)} but
114+
* some engines might want to handle this differently for better performance.
115+
* <p>
116+
* @param value per-document value, typically a String, Long, or Double
117+
*/
118+
public void setNextAggregationValue(Object value) {
119+
this.value = value;
120+
}
121+
122+
public Number get_score() {
123+
try {
124+
return scorer == null ? 0.0 : scorer.score();
125+
} catch (IOException e) {
126+
throw new ElasticsearchException("couldn't lookup score", e);
127+
}
128+
}
129+
130+
public Object get_value() {
131+
return value;
132+
}
133+
134+
/**
135+
* Return the result as a long. This is used by aggregation scripts over long fields.
136+
*/
137+
public long runAsLong() {
138+
return ((Number) execute()).longValue();
139+
}
140+
141+
public double runAsDouble() {
142+
return ((Number) execute()).doubleValue();
143+
}
144+
145+
public abstract Object execute();
146+
147+
/**
148+
* A factory to construct {@link AggregationScript} instances.
149+
*/
150+
public interface LeafFactory {
151+
AggregationScript newInstance(LeafReaderContext ctx) throws IOException;
152+
153+
/**
154+
* Return {@code true} if the script needs {@code _score} calculated, or {@code false} otherwise.
155+
*/
156+
boolean needs_score();
157+
}
158+
159+
/**
160+
* A factory to construct stateful {@link AggregationScript} factories for a specific index.
161+
*/
162+
public interface Factory {
163+
LeafFactory newFactory(Map<String, Object> params, SearchLookup lookup);
164+
}
165+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ public class ScriptModule {
4141
static {
4242
CORE_CONTEXTS = Stream.of(
4343
SearchScript.CONTEXT,
44-
SearchScript.AGGS_CONTEXT,
44+
AggregationScript.CONTEXT,
4545
ScoreScript.CONTEXT,
4646
SearchScript.SCRIPT_SORT_CONTEXT,
4747
TermsSetQueryScript.CONTEXT,

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

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@
3838
* <li>Construct a {@link LeafFactory} for a an index using {@link Factory#newFactory(Map, SearchLookup)}</li>
3939
* <li>Construct a {@link SearchScript} for a Lucene segment using {@link LeafFactory#newInstance(LeafReaderContext)}</li>
4040
* <li>Call {@link #setDocument(int)} to indicate which document in the segment the script should be run for next</li>
41-
* <li>Call one of the {@code run} methods: {@link #run()}, {@link #runAsDouble()}, or {@link #runAsLong()}</li>
41+
* <li>Call one of the {@code run} methods: {@link #run()} or {@link #runAsDouble()}</li>
4242
* </ol>
4343
*/
4444
public abstract class SearchScript implements ScorerAware {
@@ -114,10 +114,6 @@ public void setNextAggregationValue(Object value) {
114114

115115
public void setNextVar(String field, Object value) {}
116116

117-
/** Return the result as a long. This is used by aggregation scripts over long fields. */
118-
public long runAsLong() {
119-
throw new UnsupportedOperationException("runAsLong is not implemented");
120-
}
121117

122118
public Object run() {
123119
return runAsDouble();
@@ -144,7 +140,6 @@ public interface Factory {
144140
/** The context used to compile {@link SearchScript} factories. */
145141
public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("search", Factory.class);
146142
// TODO: remove these contexts when it has its own interface
147-
public static final ScriptContext<Factory> AGGS_CONTEXT = new ScriptContext<>("aggs", Factory.class);
148143
// Can return a double. (For ScriptSortType#NUMBER only, for ScriptSortType#STRING normal CONTEXT should be used)
149144
public static final ScriptContext<Factory> SCRIPT_SORT_CONTEXT = new ScriptContext<>("sort", Factory.class);
150145
}

0 commit comments

Comments
 (0)