Skip to content

Commit

Permalink
Migrate scripted metric aggregation scripts to ScriptContext design e…
Browse files Browse the repository at this point in the history
  • Loading branch information
rationull committed Apr 23, 2018
1 parent 6a4c5f3 commit 2710d5a
Show file tree
Hide file tree
Showing 10 changed files with 527 additions and 46 deletions.
@@ -0,0 +1,113 @@
/*
* Licensed to Elasticsearch under one or more contributor
* license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright
* ownership. Elasticsearch licenses this file to you under
* the Apache License, Version 2.0 (the "License"); you may
* not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/

package org.elasticsearch.painless;

import org.elasticsearch.painless.spi.Whitelist;
import org.elasticsearch.script.MetricAggScripts;
import org.elasticsearch.script.ScriptContext;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class MetricAggScriptsTests extends ScriptTestCase {
@Override
protected Map<ScriptContext<?>, List<Whitelist>> scriptContexts() {
Map<ScriptContext<?>, List<Whitelist>> contexts = new HashMap<>();
contexts.put(MetricAggScripts.InitScript.CONTEXT, Whitelist.BASE_WHITELISTS);
contexts.put(MetricAggScripts.MapScript.CONTEXT, Whitelist.BASE_WHITELISTS);
contexts.put(MetricAggScripts.CombineScript.CONTEXT, Whitelist.BASE_WHITELISTS);
contexts.put(MetricAggScripts.ReduceScript.CONTEXT, Whitelist.BASE_WHITELISTS);
return contexts;
}

public void testInitBasic() {
MetricAggScripts.InitScript.Factory factory = scriptEngine.compile("test",
"agg.testField = params.initialVal", MetricAggScripts.InitScript.CONTEXT, Collections.emptyMap());

Map<String, Object> params = new HashMap<>();
Map<String, Object> agg = new HashMap<>();

params.put("initialVal", 10);

MetricAggScripts.InitScript script = factory.newInstance(params, agg);
script.execute();

assert(agg.containsKey("testField"));
assertEquals(10, agg.get("testField"));
}

public void testMapBasic() {
MetricAggScripts.MapScript.Factory factory = scriptEngine.compile("test",
"agg.testField = 2*_score", MetricAggScripts.MapScript.CONTEXT, Collections.emptyMap());

Map<String, Object> params = new HashMap<>();
Map<String, Object> agg = new HashMap<>();
double _score = 0.5;

MetricAggScripts.MapScript.LeafFactory leafFactory = factory.newFactory(params, agg, null);
MetricAggScripts.MapScript script = leafFactory.newInstance(null);

script.execute(_score);

assert(agg.containsKey("testField"));
assertEquals(1.0, agg.get("testField"));
}

public void testCombineBasic() {
MetricAggScripts.CombineScript.Factory factory = scriptEngine.compile("test",
"agg.testField = params.initialVal; return agg.testField + params.inc", MetricAggScripts.CombineScript.CONTEXT,
Collections.emptyMap());

Map<String, Object> params = new HashMap<>();
Map<String, Object> agg = new HashMap<>();

params.put("initialVal", 10);
params.put("inc", 2);

MetricAggScripts.CombineScript script = factory.newInstance(params, agg);
Object res = script.execute();

assert(agg.containsKey("testField"));
assertEquals(10, agg.get("testField"));
assertEquals(12, res);
}

public void testReduceBasic() {
MetricAggScripts.ReduceScript.Factory factory = scriptEngine.compile("test",
"aggs[0].testField + aggs[1].testField", MetricAggScripts.ReduceScript.CONTEXT, Collections.emptyMap());

Map<String, Object> params = new HashMap<>();
List<Object> aggs = new ArrayList<>();

Map<String, Object> agg1 = new HashMap<>(), agg2 = new HashMap<>();
agg1.put("testField", 1);
agg2.put("testField", 2);

aggs.add(agg1);
aggs.add(agg2);

MetricAggScripts.ReduceScript script = factory.newInstance(params, aggs);
Object res = script.execute();
assertEquals(3, res);
}
}
140 changes: 140 additions & 0 deletions server/src/main/java/org/elasticsearch/script/MetricAggScripts.java
@@ -0,0 +1,140 @@
/*
* Licensed to Elasticsearch under one or more contributor
* license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright
* ownership. Elasticsearch licenses this file to you under
* the Apache License, Version 2.0 (the "License"); you may
* not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/

package org.elasticsearch.script;

import org.apache.lucene.index.LeafReaderContext;
import org.elasticsearch.index.fielddata.ScriptDocValues;
import org.elasticsearch.search.lookup.LeafSearchLookup;
import org.elasticsearch.search.lookup.SearchLookup;

import java.util.List;
import java.util.Map;

public class MetricAggScripts {
private abstract static class ParamsAndAggBase {
private final Map<String, Object> params;
private final Object agg;

ParamsAndAggBase(Map<String, Object> params, Object agg) {
this.params = params;
this.agg = agg;
}

public Map<String, Object> getParams() {
return params;
}

public Object getAgg() {
return agg;
}
}

public abstract static class InitScript extends ParamsAndAggBase {
public InitScript(Map<String, Object> params, Object agg) {
super(params, agg);
}

public abstract void execute();

public interface Factory {
InitScript newInstance(Map<String, Object> params, Object agg);
}

public static String[] PARAMETERS = {};
public static ScriptContext<Factory> CONTEXT = new ScriptContext<>("aggs_init", Factory.class);
}

public abstract static class MapScript extends ParamsAndAggBase {
private final LeafSearchLookup leafLookup;

public MapScript(Map<String, Object> params, Object agg, SearchLookup lookup, LeafReaderContext leafContext) {
super(params, agg);

this.leafLookup = leafContext == null ? null : lookup.getLeafSearchLookup(leafContext);
}

// Return the doc as a map (instead of LeafDocLookup) in order to abide by type whitelisting rules for
// Painless scripts.
public Map<String, ScriptDocValues<?>> getDoc() {
return leafLookup == null ? null : leafLookup.doc();
}

public void setDocument(int docId) {
if (leafLookup != null) {
leafLookup.setDocument(docId);
}
}

public abstract void execute(double _score);

public interface LeafFactory {
MapScript newInstance(LeafReaderContext ctx);
}

public interface Factory {
LeafFactory newFactory(Map<String, Object> params, Object agg, SearchLookup lookup);
}

public static String[] PARAMETERS = new String[] {"_score"};
public static ScriptContext<Factory> CONTEXT = new ScriptContext<>("aggs_map", Factory.class);
}

public abstract static class CombineScript extends ParamsAndAggBase {
public CombineScript(Map<String, Object> params, Object agg) {
super(params, agg);
}

public abstract Object execute();

public interface Factory {
CombineScript newInstance(Map<String, Object> params, Object agg);
}

public static String[] PARAMETERS = {};
public static ScriptContext<Factory> CONTEXT = new ScriptContext<>("aggs_combine", Factory.class);
}

public abstract static class ReduceScript {
private final Map<String, Object> params;
private final List<Object> aggs;

public ReduceScript(Map<String, Object> params, List<Object> aggs) {
this.params = params;
this.aggs = aggs;
}

public Map<String, Object> getParams() {
return params;
}

public List<Object> getAggs() {
return aggs;
}

public abstract Object execute();

public interface Factory {
ReduceScript newInstance(Map<String, Object> params, List<Object> aggs);
}

public static String[] PARAMETERS = {};
public static ScriptContext<Factory> CONTEXT = new ScriptContext<>("aggs_reduce", Factory.class);
}
}
Expand Up @@ -48,7 +48,11 @@ public class ScriptModule {
FilterScript.CONTEXT,
SimilarityScript.CONTEXT,
SimilarityWeightScript.CONTEXT,
TemplateScript.CONTEXT
TemplateScript.CONTEXT,
MetricAggScripts.InitScript.CONTEXT,
MetricAggScripts.MapScript.CONTEXT,
MetricAggScripts.CombineScript.CONTEXT,
MetricAggScripts.ReduceScript.CONTEXT
).collect(Collectors.toMap(c -> c.name, Function.identity()));
}

Expand Down
Expand Up @@ -22,7 +22,7 @@
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.script.ExecutableScript;
import org.elasticsearch.script.MetricAggScripts;
import org.elasticsearch.script.Script;
import org.elasticsearch.search.aggregations.InternalAggregation;
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
Expand Down Expand Up @@ -89,15 +89,16 @@ public InternalAggregation doReduce(List<InternalAggregation> aggregations, Redu
InternalScriptedMetric firstAggregation = ((InternalScriptedMetric) aggregations.get(0));
List<Object> aggregation;
if (firstAggregation.reduceScript != null && reduceContext.isFinalReduce()) {
Map<String, Object> vars = new HashMap<>();
vars.put("_aggs", aggregationObjects);
Map<String, Object> params = new HashMap<>();
if (firstAggregation.reduceScript.getParams() != null) {
vars.putAll(firstAggregation.reduceScript.getParams());
params.putAll(firstAggregation.reduceScript.getParams());
}
ExecutableScript.Factory factory = reduceContext.scriptService().compile(
firstAggregation.reduceScript, ExecutableScript.AGGS_CONTEXT);
ExecutableScript script = factory.newInstance(vars);
aggregation = Collections.singletonList(script.run());
params.put("_aggs", aggregationObjects);

MetricAggScripts.ReduceScript.Factory factory = reduceContext.scriptService().compile(
firstAggregation.reduceScript, MetricAggScripts.ReduceScript.CONTEXT);
MetricAggScripts.ReduceScript script = factory.newInstance(params, aggregationObjects);
aggregation = Collections.singletonList(script.execute());
} else if (reduceContext.isFinalReduce()) {
aggregation = Collections.singletonList(aggregationObjects);
} else {
Expand Down
Expand Up @@ -26,12 +26,10 @@
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.script.ExecutableScript;
import org.elasticsearch.script.MetricAggScripts;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.SearchScript;
import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.AggregatorFactories;
import org.elasticsearch.search.aggregations.AggregatorFactories.Builder;
import org.elasticsearch.search.aggregations.AggregatorFactory;
import org.elasticsearch.search.internal.SearchContext;
Expand Down Expand Up @@ -203,30 +201,32 @@ protected ScriptedMetricAggregatorFactory doBuild(SearchContext context, Aggrega
// Extract params from scripts and pass them along to ScriptedMetricAggregatorFactory, since it won't have
// access to them for the scripts it's given precompiled.

ExecutableScript.Factory executableInitScript;
MetricAggScripts.InitScript.Factory compiledInitScript;
Map<String, Object> initScriptParams;
if (initScript != null) {
executableInitScript = queryShardContext.getScriptService().compile(initScript, ExecutableScript.AGGS_CONTEXT);
compiledInitScript = queryShardContext.getScriptService().compile(initScript, MetricAggScripts.InitScript.CONTEXT);
initScriptParams = initScript.getParams();
} else {
executableInitScript = p -> null;
compiledInitScript = (p, a) -> null;
initScriptParams = Collections.emptyMap();
}

SearchScript.Factory searchMapScript = queryShardContext.getScriptService().compile(mapScript, SearchScript.AGGS_CONTEXT);
MetricAggScripts.MapScript.Factory compiledMapScript = queryShardContext.getScriptService().compile(mapScript,
MetricAggScripts.MapScript.CONTEXT);
Map<String, Object> mapScriptParams = mapScript.getParams();

ExecutableScript.Factory executableCombineScript;
MetricAggScripts.CombineScript.Factory compiledCombineScript;
Map<String, Object> combineScriptParams;
if (combineScript != null) {
executableCombineScript = queryShardContext.getScriptService().compile(combineScript, ExecutableScript.AGGS_CONTEXT);
compiledCombineScript = queryShardContext.getScriptService().compile(combineScript,
MetricAggScripts.CombineScript.CONTEXT);
combineScriptParams = combineScript.getParams();
} else {
executableCombineScript = p -> null;
compiledCombineScript = (p, a) -> null;
combineScriptParams = Collections.emptyMap();
}
return new ScriptedMetricAggregatorFactory(name, searchMapScript, mapScriptParams, executableInitScript, initScriptParams,
executableCombineScript, combineScriptParams, reduceScript,
return new ScriptedMetricAggregatorFactory(name, compiledMapScript, mapScriptParams, compiledInitScript,
initScriptParams, compiledCombineScript, combineScriptParams, reduceScript,
params, queryShardContext.lookup(), context, parent, subfactoriesBuilder, metaData);
}

Expand Down

0 comments on commit 2710d5a

Please sign in to comment.