Skip to content

Commit 6fa7016

Browse files
original-brownbearoriginal-brownbear
authored
SCRIPTING: Move Aggregation Scripts to their own context (#32068)
* SCRIPTING: Move Aggregation Scripts to their own context
1 parent 6ca24e1 commit 6fa7016

File tree

10 files changed

+254
-53
lines changed

10 files changed

+254
-53
lines changed

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

+50-5
Original file line numberDiff line numberDiff line change
@@ -34,10 +34,12 @@
3434
import org.elasticsearch.common.settings.Settings;
3535
import org.elasticsearch.index.fielddata.IndexFieldData;
3636
import org.elasticsearch.index.fielddata.IndexNumericFieldData;
37-
import org.elasticsearch.index.mapper.GeoPointFieldMapper.GeoPointFieldType;
3837
import org.elasticsearch.index.mapper.DateFieldMapper;
38+
import org.elasticsearch.index.mapper.GeoPointFieldMapper.GeoPointFieldType;
3939
import org.elasticsearch.index.mapper.MappedFieldType;
4040
import org.elasticsearch.index.mapper.MapperService;
41+
import org.elasticsearch.script.BucketAggregationScript;
42+
import org.elasticsearch.script.BucketAggregationSelectorScript;
4143
import org.elasticsearch.script.ClassPermission;
4244
import org.elasticsearch.script.ExecutableScript;
4345
import org.elasticsearch.script.FilterScript;
@@ -54,6 +56,7 @@
5456
import java.security.PrivilegedAction;
5557
import java.text.ParseException;
5658
import java.util.ArrayList;
59+
import java.util.HashMap;
5760
import java.util.List;
5861
import java.util.Map;
5962

@@ -112,6 +115,17 @@ protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundE
112115
} else if (context.instanceClazz.equals(ExecutableScript.class)) {
113116
ExecutableScript.Factory factory = (p) -> new ExpressionExecutableScript(expr, p);
114117
return context.factoryClazz.cast(factory);
118+
} else if (context.instanceClazz.equals(BucketAggregationScript.class)) {
119+
return context.factoryClazz.cast(newBucketAggregationScriptFactory(expr));
120+
} else if (context.instanceClazz.equals(BucketAggregationSelectorScript.class)) {
121+
BucketAggregationScript.Factory factory = newBucketAggregationScriptFactory(expr);
122+
BucketAggregationSelectorScript.Factory wrappedFactory = parameters -> new BucketAggregationSelectorScript(parameters) {
123+
@Override
124+
public boolean execute() {
125+
return factory.newInstance(getParams()).execute() == 1.0;
126+
}
127+
};
128+
return context.factoryClazz.cast(wrappedFactory);
115129
} else if (context.instanceClazz.equals(FilterScript.class)) {
116130
FilterScript.Factory factory = (p, lookup) -> newFilterScript(expr, lookup, p);
117131
return context.factoryClazz.cast(factory);
@@ -122,6 +136,37 @@ protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundE
122136
throw new IllegalArgumentException("expression engine does not know how to handle script context [" + context.name + "]");
123137
}
124138

139+
private static BucketAggregationScript.Factory newBucketAggregationScriptFactory(Expression expr) {
140+
return parameters -> {
141+
ReplaceableConstDoubleValues[] functionValuesArray =
142+
new ReplaceableConstDoubleValues[expr.variables.length];
143+
Map<String, ReplaceableConstDoubleValues> functionValuesMap = new HashMap<>();
144+
for (int i = 0; i < expr.variables.length; ++i) {
145+
functionValuesArray[i] = new ReplaceableConstDoubleValues();
146+
functionValuesMap.put(expr.variables[i], functionValuesArray[i]);
147+
}
148+
return new BucketAggregationScript(parameters) {
149+
@Override
150+
public double execute() {
151+
getParams().forEach((name, value) -> {
152+
ReplaceableConstDoubleValues placeholder = functionValuesMap.get(name);
153+
if (placeholder == null) {
154+
throw new IllegalArgumentException("Error using " + expr + ". " +
155+
"The variable [" + name + "] does not exist in the executable expressions script.");
156+
} else if (value instanceof Number == false) {
157+
throw new IllegalArgumentException("Error using " + expr + ". " +
158+
"Executable expressions scripts can only process numbers." +
159+
" The variable [" + name + "] is not a number.");
160+
} else {
161+
placeholder.setValue(((Number) value).doubleValue());
162+
}
163+
});
164+
return expr.evaluate(functionValuesArray);
165+
}
166+
};
167+
};
168+
}
169+
125170
private SearchScript.LeafFactory newSearchScript(Expression expr, SearchLookup lookup, @Nullable Map<String, Object> vars) {
126171
MapperService mapper = lookup.doc().mapperService();
127172
// NOTE: if we need to do anything complicated with bindings in the future, we can just extend Bindings,
@@ -267,7 +312,7 @@ public void setDocument(int docid) {
267312
};
268313
};
269314
}
270-
315+
271316
private ScoreScript.LeafFactory newScoreScript(Expression expr, SearchLookup lookup, @Nullable Map<String, Object> vars) {
272317
SearchScript.LeafFactory searchLeafFactory = newSearchScript(expr, lookup, vars);
273318
return new ScoreScript.LeafFactory() {
@@ -284,17 +329,17 @@ public ScoreScript newInstance(LeafReaderContext ctx) throws IOException {
284329
public double execute() {
285330
return script.runAsDouble();
286331
}
287-
332+
288333
@Override
289334
public void setDocument(int docid) {
290335
script.setDocument(docid);
291336
}
292-
337+
293338
@Override
294339
public void setScorer(Scorer scorer) {
295340
script.setScorer(scorer);
296341
}
297-
342+
298343
@Override
299344
public double get_score() {
300345
return script.getScore();

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

+54
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
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.Map;
23+
24+
/**
25+
* A script used in bucket aggregations that returns a {@code double} value.
26+
*/
27+
public abstract class BucketAggregationScript {
28+
29+
public static final String[] PARAMETERS = {};
30+
31+
public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("bucket_aggregation", Factory.class);
32+
33+
/**
34+
* The generic runtime parameters for the script.
35+
*/
36+
private final Map<String, Object> params;
37+
38+
public BucketAggregationScript(Map<String, Object> params) {
39+
this.params = params;
40+
}
41+
42+
/**
43+
* Return the parameters for this script.
44+
*/
45+
public Map<String, Object> getParams() {
46+
return params;
47+
}
48+
49+
public abstract double execute();
50+
51+
public interface Factory {
52+
BucketAggregationScript newInstance(Map<String, Object> params);
53+
}
54+
}

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

+54
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
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.Map;
23+
24+
/**
25+
* A script used in bucket aggregations that returns a {@code boolean} value.
26+
*/
27+
public abstract class BucketAggregationSelectorScript {
28+
29+
public static final String[] PARAMETERS = {};
30+
31+
public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("aggregation_selector", Factory.class);
32+
33+
/**
34+
* The generic runtime parameters for the script.
35+
*/
36+
private final Map<String, Object> params;
37+
38+
public BucketAggregationSelectorScript(Map<String, Object> params) {
39+
this.params = params;
40+
}
41+
42+
/**
43+
* Return the parameters for this script.
44+
*/
45+
public Map<String, Object> getParams() {
46+
return params;
47+
}
48+
49+
public abstract boolean execute();
50+
51+
public interface Factory {
52+
BucketAggregationSelectorScript newInstance(Map<String, Object> params);
53+
}
54+
}

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

-1
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,5 @@ interface Factory {
4848
ScriptContext<Factory> CONTEXT = new ScriptContext<>("executable", Factory.class);
4949

5050
// TODO: remove these once each has its own script interface
51-
ScriptContext<Factory> AGGS_CONTEXT = new ScriptContext<>("aggs_executable", Factory.class);
5251
ScriptContext<Factory> UPDATE_CONTEXT = new ScriptContext<>("update", Factory.class);
5352
}

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

+3-1
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,9 @@ public class ScriptModule {
4646
SearchScript.SCRIPT_SORT_CONTEXT,
4747
SearchScript.TERMS_SET_QUERY_CONTEXT,
4848
ExecutableScript.CONTEXT,
49-
ExecutableScript.AGGS_CONTEXT,
49+
BucketAggregationScript.CONTEXT,
50+
BucketAggregationSelectorScript.CONTEXT,
51+
SignificantTermsHeuristicScoreScript.CONTEXT,
5052
ExecutableScript.UPDATE_CONTEXT,
5153
IngestScript.CONTEXT,
5254
FilterScript.CONTEXT,

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

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
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.Map;
23+
24+
/**
25+
* A script used in significant terms heuristic scoring.
26+
*/
27+
public abstract class SignificantTermsHeuristicScoreScript {
28+
29+
public static final String[] PARAMETERS = { "params" };
30+
31+
public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("script_heuristic", Factory.class);
32+
33+
public abstract double execute(Map<String, Object> params);
34+
35+
public interface Factory {
36+
SignificantTermsHeuristicScoreScript newInstance();
37+
}
38+
}

‎server/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/ScriptHeuristic.java

+16-12
Original file line numberDiff line numberDiff line change
@@ -28,12 +28,14 @@
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.ExecutableScript;
3231
import org.elasticsearch.script.Script;
32+
import org.elasticsearch.script.SignificantTermsHeuristicScoreScript;
3333
import org.elasticsearch.search.aggregations.InternalAggregation;
3434
import org.elasticsearch.search.internal.SearchContext;
3535

3636
import java.io.IOException;
37+
import java.util.HashMap;
38+
import java.util.Map;
3739
import java.util.Objects;
3840

3941
public class ScriptHeuristic extends SignificanceHeuristic {
@@ -48,19 +50,21 @@ static class ExecutableScriptHeuristic extends ScriptHeuristic {
4850
private final LongAccessor supersetSizeHolder;
4951
private final LongAccessor subsetDfHolder;
5052
private final LongAccessor supersetDfHolder;
51-
private final ExecutableScript executableScript;
53+
private final SignificantTermsHeuristicScoreScript executableScript;
54+
private final Map<String, Object> params = new HashMap<>();
5255

53-
ExecutableScriptHeuristic(Script script, ExecutableScript executableScript){
56+
ExecutableScriptHeuristic(Script script, SignificantTermsHeuristicScoreScript executableScript) {
5457
super(script);
5558
subsetSizeHolder = new LongAccessor();
5659
supersetSizeHolder = new LongAccessor();
5760
subsetDfHolder = new LongAccessor();
5861
supersetDfHolder = new LongAccessor();
5962
this.executableScript = executableScript;
60-
executableScript.setNextVar("_subset_freq", subsetDfHolder);
61-
executableScript.setNextVar("_subset_size", subsetSizeHolder);
62-
executableScript.setNextVar("_superset_freq", supersetDfHolder);
63-
executableScript.setNextVar("_superset_size", supersetSizeHolder);
63+
params.putAll(script.getParams());
64+
params.put("_subset_freq", subsetDfHolder);
65+
params.put("_subset_size", subsetSizeHolder);
66+
params.put("_superset_freq", supersetDfHolder);
67+
params.put("_superset_size", supersetSizeHolder);
6468
}
6569

6670
@Override
@@ -69,7 +73,7 @@ public double getScore(long subsetFreq, long subsetSize, long supersetFreq, long
6973
supersetSizeHolder.value = supersetSize;
7074
subsetDfHolder.value = subsetFreq;
7175
supersetDfHolder.value = supersetFreq;
72-
return ((Number) executableScript.run()).doubleValue();
76+
return executableScript.execute(params);
7377
}
7478
}
7579

@@ -91,15 +95,15 @@ public void writeTo(StreamOutput out) throws IOException {
9195

9296
@Override
9397
public SignificanceHeuristic rewrite(InternalAggregation.ReduceContext context) {
94-
ExecutableScript.Factory factory = context.scriptService().compile(script, ExecutableScript.AGGS_CONTEXT);
95-
return new ExecutableScriptHeuristic(script, factory.newInstance(script.getParams()));
98+
SignificantTermsHeuristicScoreScript.Factory factory = context.scriptService().compile(script, SignificantTermsHeuristicScoreScript.CONTEXT);
99+
return new ExecutableScriptHeuristic(script, factory.newInstance());
96100
}
97101

98102
@Override
99103
public SignificanceHeuristic rewrite(SearchContext context) {
100104
QueryShardContext shardContext = context.getQueryShardContext();
101-
ExecutableScript.Factory compiledScript = shardContext.getScriptService().compile(script, ExecutableScript.AGGS_CONTEXT);
102-
return new ExecutableScriptHeuristic(script, compiledScript.newInstance(script.getParams()));
105+
SignificantTermsHeuristicScoreScript.Factory compiledScript = shardContext.getScriptService().compile(script, SignificantTermsHeuristicScoreScript.CONTEXT);
106+
return new ExecutableScriptHeuristic(script, compiledScript.newInstance());
103107
}
104108

105109

0 commit comments

Comments
 (0)