SCRIPTING: Move Aggregation Script Context to its own class (#33820)

* SCRIPTING: Move Aggregation Script Context to its own class
This commit is contained in:
Armin Braun 2018-10-15 17:28:05 +01:00 committed by GitHub
parent 511526250b
commit ebca27371c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
18 changed files with 384 additions and 76 deletions

View File

@ -0,0 +1,93 @@
/*
* 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.expression;
import java.io.IOException;
import org.apache.lucene.expressions.Bindings;
import org.apache.lucene.expressions.Expression;
import org.apache.lucene.expressions.SimpleBindings;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.DoubleValues;
import org.apache.lucene.search.DoubleValuesSource;
import org.elasticsearch.script.AggregationScript;
import org.elasticsearch.script.GeneralScriptException;
/**
* A bridge to evaluate an {@link Expression} against {@link Bindings} in the context
* of a {@link AggregationScript}.
*/
class ExpressionAggregationScript implements AggregationScript.LeafFactory {
final Expression exprScript;
final SimpleBindings bindings;
final DoubleValuesSource source;
final ReplaceableConstDoubleValueSource specialValue; // _value
ExpressionAggregationScript(Expression e, SimpleBindings b, ReplaceableConstDoubleValueSource v) {
exprScript = e;
bindings = b;
source = exprScript.getDoubleValuesSource(bindings);
specialValue = v;
}
@Override
public AggregationScript newInstance(final LeafReaderContext leaf) throws IOException {
return new AggregationScript() {
// Fake the scorer until setScorer is called.
DoubleValues values = source.getValues(leaf, null);
@Override
public Object execute() {
try {
return values.doubleValue();
} catch (Exception exception) {
throw new GeneralScriptException("Error evaluating " + exprScript, exception);
}
}
@Override
public void setDocument(int d) {
try {
values.advanceExact(d);
} catch (IOException e) {
throw new IllegalStateException("Can't advance to doc using " + exprScript, e);
}
}
@Override
public void setNextAggregationValue(Object value) {
// _value isn't used in script if specialValue == null
if (specialValue != null) {
if (value instanceof Number) {
specialValue.setValue(((Number)value).doubleValue());
} else {
throw new GeneralScriptException("Cannot use expression with text variable using " + exprScript);
}
}
}
};
}
@Override
public boolean needs_score() {
return false;
}
}

View File

@ -37,6 +37,7 @@ import org.elasticsearch.index.fielddata.IndexNumericFieldData;
import org.elasticsearch.index.mapper.DateFieldMapper;
import org.elasticsearch.index.mapper.GeoPointFieldMapper.GeoPointFieldType;
import org.elasticsearch.index.mapper.MappedFieldType;
import org.elasticsearch.script.AggregationScript;
import org.elasticsearch.script.BucketAggregationScript;
import org.elasticsearch.script.BucketAggregationSelectorScript;
import org.elasticsearch.script.ClassPermission;
@ -131,6 +132,9 @@ public class ExpressionScriptEngine extends AbstractComponent implements ScriptE
} else if (context.instanceClazz.equals(TermsSetQueryScript.class)) {
TermsSetQueryScript.Factory factory = (p, lookup) -> newTermsSetQueryScript(expr, lookup, p);
return context.factoryClazz.cast(factory);
} else if (context.instanceClazz.equals(AggregationScript.class)) {
AggregationScript.Factory factory = (p, lookup) -> newAggregationScript(expr, lookup, p);
return context.factoryClazz.cast(factory);
}
throw new IllegalArgumentException("expression engine does not know how to handle script context [" + context.name + "]");
}
@ -224,6 +228,37 @@ public class ExpressionScriptEngine extends AbstractComponent implements ScriptE
return new ExpressionTermSetQueryScript(expr, bindings);
}
private AggregationScript.LeafFactory newAggregationScript(Expression expr, SearchLookup lookup,
@Nullable Map<String, Object> vars) {
// NOTE: if we need to do anything complicated with bindings in the future, we can just extend Bindings,
// instead of complicating SimpleBindings (which should stay simple)
SimpleBindings bindings = new SimpleBindings();
ReplaceableConstDoubleValueSource specialValue = null;
for (String variable : expr.variables) {
try {
if (variable.equals("_value")) {
specialValue = new ReplaceableConstDoubleValueSource();
bindings.add("_value", specialValue);
// noop: _value is special for aggregations, and is handled in ExpressionScriptBindings
// TODO: if some uses it in a scoring expression, they will get a nasty failure when evaluating...need a
// way to know this is for aggregations and so _value is ok to have...
} else if (vars != null && vars.containsKey(variable)) {
bindFromParams(vars, bindings, variable);
} else {
// delegate valuesource creation based on field's type
// there are three types of "fields" to expressions, and each one has a different "api" of variables and methods.
final ValueSource valueSource = getDocValueSource(variable, lookup);
bindings.add(variable, valueSource.asDoubleValuesSource());
}
} catch (Exception e) {
// we defer "binding" of variables until here: give context for that variable
throw convertToScriptException("link error", expr.sourceText, variable, e);
}
}
return new ExpressionAggregationScript(expr, bindings, specialValue);
}
/**
* This is a hack for filter scripts, which must return booleans instead of doubles as expression do.
* See https://github.com/elastic/elasticsearch/issues/26429.

View File

@ -75,9 +75,6 @@ class ExpressionSearchScript implements SearchScript.LeafFactory {
@Override
public Object run() { return Double.valueOf(runAsDouble()); }
@Override
public long runAsLong() { return (long)runAsDouble(); }
@Override
public double runAsDouble() {
try {

View File

@ -109,8 +109,4 @@ final class ScriptImpl extends SearchScript {
return ((Number)run()).doubleValue();
}
@Override
public long runAsLong() {
return ((Number)run()).longValue();
}
}

View File

@ -0,0 +1,165 @@
/*
* 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 java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.Scorable;
import org.elasticsearch.ElasticsearchException;
import org.elasticsearch.common.lucene.ScorerAware;
import org.elasticsearch.index.fielddata.ScriptDocValues;
import org.elasticsearch.search.lookup.LeafSearchLookup;
import org.elasticsearch.search.lookup.SearchLookup;
public abstract class AggregationScript implements ScorerAware {
public static final String[] PARAMETERS = {};
public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("aggs", Factory.class);
private static final Map<String, String> DEPRECATIONS;
static {
Map<String, String> deprecations = new HashMap<>();
deprecations.put(
"doc",
"Accessing variable [doc] via [params.doc] from within an aggregation-script " +
"is deprecated in favor of directly accessing [doc]."
);
deprecations.put(
"_doc",
"Accessing variable [doc] via [params._doc] from within an aggregation-script " +
"is deprecated in favor of directly accessing [doc]."
);
DEPRECATIONS = Collections.unmodifiableMap(deprecations);
}
/**
* The generic runtime parameters for the script.
*/
private final Map<String, Object> params;
/**
* A leaf lookup for the bound segment this script will operate on.
*/
private final LeafSearchLookup leafLookup;
/**
* A scorer that will return the score for the current document when the script is run.
*/
protected Scorable scorer;
private Object value;
public AggregationScript(Map<String, Object> params, SearchLookup lookup, LeafReaderContext leafContext) {
this.params = new ParameterMap(new HashMap<>(params), DEPRECATIONS);
this.leafLookup = lookup.getLeafSearchLookup(leafContext);
this.params.putAll(leafLookup.asMap());
}
protected AggregationScript() {
params = null;
leafLookup = null;
}
/**
* Return the parameters for this script.
*/
public Map<String, Object> getParams() {
return params;
}
/**
* The doc lookup for the Lucene segment this script was created for.
*/
public Map<String, ScriptDocValues<?>> getDoc() {
return leafLookup.doc();
}
/**
* Set the current document to run the script on next.
*/
public void setDocument(int docid) {
leafLookup.setDocument(docid);
}
@Override
public void setScorer(Scorable scorer) {
this.scorer = scorer;
}
/**
* Sets per-document aggregation {@code _value}.
* <p>
* The default implementation just calls {@code setNextVar("_value", value)} but
* some engines might want to handle this differently for better performance.
* <p>
* @param value per-document value, typically a String, Long, or Double
*/
public void setNextAggregationValue(Object value) {
this.value = value;
}
public Number get_score() {
try {
return scorer == null ? 0.0 : scorer.score();
} catch (IOException e) {
throw new ElasticsearchException("couldn't lookup score", e);
}
}
public Object get_value() {
return value;
}
/**
* Return the result as a long. This is used by aggregation scripts over long fields.
*/
public long runAsLong() {
return ((Number) execute()).longValue();
}
public double runAsDouble() {
return ((Number) execute()).doubleValue();
}
public abstract Object execute();
/**
* A factory to construct {@link AggregationScript} instances.
*/
public interface LeafFactory {
AggregationScript newInstance(LeafReaderContext ctx) throws IOException;
/**
* Return {@code true} if the script needs {@code _score} calculated, or {@code false} otherwise.
*/
boolean needs_score();
}
/**
* A factory to construct stateful {@link AggregationScript} factories for a specific index.
*/
public interface Factory {
LeafFactory newFactory(Map<String, Object> params, SearchLookup lookup);
}
}

View File

@ -41,7 +41,7 @@ public class ScriptModule {
static {
CORE_CONTEXTS = Stream.of(
SearchScript.CONTEXT,
SearchScript.AGGS_CONTEXT,
AggregationScript.CONTEXT,
ScoreScript.CONTEXT,
SearchScript.SCRIPT_SORT_CONTEXT,
TermsSetQueryScript.CONTEXT,

View File

@ -38,7 +38,7 @@ import java.util.Map;
* <li>Construct a {@link LeafFactory} for a an index using {@link Factory#newFactory(Map, SearchLookup)}</li>
* <li>Construct a {@link SearchScript} for a Lucene segment using {@link LeafFactory#newInstance(LeafReaderContext)}</li>
* <li>Call {@link #setDocument(int)} to indicate which document in the segment the script should be run for next</li>
* <li>Call one of the {@code run} methods: {@link #run()}, {@link #runAsDouble()}, or {@link #runAsLong()}</li>
* <li>Call one of the {@code run} methods: {@link #run()} or {@link #runAsDouble()}</li>
* </ol>
*/
public abstract class SearchScript implements ScorerAware {
@ -114,10 +114,6 @@ public abstract class SearchScript implements ScorerAware {
public void setNextVar(String field, Object value) {}
/** Return the result as a long. This is used by aggregation scripts over long fields. */
public long runAsLong() {
throw new UnsupportedOperationException("runAsLong is not implemented");
}
public Object run() {
return runAsDouble();
@ -144,7 +140,6 @@ public abstract class SearchScript implements ScorerAware {
/** The context used to compile {@link SearchScript} factories. */
public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("search", Factory.class);
// TODO: remove these contexts when it has its own interface
public static final ScriptContext<Factory> AGGS_CONTEXT = new ScriptContext<>("aggs", Factory.class);
// Can return a double. (For ScriptSortType#NUMBER only, for ScriptSortType#STRING normal CONTEXT should be used)
public static final ScriptContext<Factory> SCRIPT_SORT_CONTEXT = new ScriptContext<>("sort", Factory.class);
}

View File

@ -42,7 +42,7 @@ import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
import org.elasticsearch.index.fielddata.SortingBinaryDocValues;
import org.elasticsearch.index.fielddata.SortingNumericDoubleValues;
import org.elasticsearch.script.SearchScript;
import org.elasticsearch.script.AggregationScript;
import org.elasticsearch.search.aggregations.support.ValuesSource.WithScript.BytesValues;
import org.elasticsearch.search.aggregations.support.values.ScriptBytesValues;
import org.elasticsearch.search.aggregations.support.values.ScriptDoubleValues;
@ -183,9 +183,9 @@ public abstract class ValuesSource {
public static class Script extends Bytes {
private final SearchScript.LeafFactory script;
private final AggregationScript.LeafFactory script;
public Script(SearchScript.LeafFactory script) {
public Script(AggregationScript.LeafFactory script) {
this.script = script;
}
@ -199,8 +199,6 @@ public abstract class ValuesSource {
return script.needs_score();
}
}
}
public abstract static class Numeric extends ValuesSource {
@ -252,9 +250,9 @@ public abstract class ValuesSource {
public static class WithScript extends Numeric {
private final Numeric delegate;
private final SearchScript.LeafFactory script;
private final AggregationScript.LeafFactory script;
public WithScript(Numeric delegate, SearchScript.LeafFactory script) {
public WithScript(Numeric delegate, AggregationScript.LeafFactory script) {
this.delegate = delegate;
this.script = script;
}
@ -287,9 +285,9 @@ public abstract class ValuesSource {
static class LongValues extends AbstractSortingNumericDocValues implements ScorerAware {
private final SortedNumericDocValues longValues;
private final SearchScript script;
private final AggregationScript script;
LongValues(SortedNumericDocValues values, SearchScript script) {
LongValues(SortedNumericDocValues values, AggregationScript script) {
this.longValues = values;
this.script = script;
}
@ -318,9 +316,9 @@ public abstract class ValuesSource {
static class DoubleValues extends SortingNumericDoubleValues implements ScorerAware {
private final SortedNumericDoubleValues doubleValues;
private final SearchScript script;
private final AggregationScript script;
DoubleValues(SortedNumericDoubleValues values, SearchScript script) {
DoubleValues(SortedNumericDoubleValues values, AggregationScript script) {
this.doubleValues = values;
this.script = script;
}
@ -377,10 +375,10 @@ public abstract class ValuesSource {
}
public static class Script extends Numeric {
private final SearchScript.LeafFactory script;
private final AggregationScript.LeafFactory script;
private final ValueType scriptValueType;
public Script(SearchScript.LeafFactory script, ValueType scriptValueType) {
public Script(AggregationScript.LeafFactory script, ValueType scriptValueType) {
this.script = script;
this.scriptValueType = scriptValueType;
}
@ -417,9 +415,9 @@ public abstract class ValuesSource {
public static class WithScript extends Bytes {
private final ValuesSource delegate;
private final SearchScript.LeafFactory script;
private final AggregationScript.LeafFactory script;
public WithScript(ValuesSource delegate, SearchScript.LeafFactory script) {
public WithScript(ValuesSource delegate, AggregationScript.LeafFactory script) {
this.delegate = delegate;
this.script = script;
}
@ -437,9 +435,9 @@ public abstract class ValuesSource {
static class BytesValues extends SortingBinaryDocValues implements ScorerAware {
private final SortedBinaryDocValues bytesValues;
private final SearchScript script;
private final AggregationScript script;
BytesValues(SortedBinaryDocValues bytesValues, SearchScript script) {
BytesValues(SortedBinaryDocValues bytesValues, AggregationScript script) {
this.bytesValues = bytesValues;
this.script = script;
}
@ -457,7 +455,7 @@ public abstract class ValuesSource {
for (int i = 0; i < count; ++i) {
final BytesRef value = bytesValues.nextValue();
script.setNextAggregationValue(value.utf8ToString());
Object run = script.run();
Object run = script.execute();
CollectionUtils.ensureNoSelfReferences(run, "ValuesSource.BytesValues script");
values[i].copyChars(run.toString());
}

View File

@ -27,8 +27,8 @@ import org.elasticsearch.index.fielddata.IndexNumericFieldData;
import org.elasticsearch.index.fielddata.IndexOrdinalsFieldData;
import org.elasticsearch.index.mapper.MappedFieldType;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.script.AggregationScript;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.SearchScript;
import org.elasticsearch.search.DocValueFormat;
import org.elasticsearch.search.aggregations.AggregationExecutionException;
import org.joda.time.DateTimeZone;
@ -113,11 +113,11 @@ public class ValuesSourceConfig<VS extends ValuesSource> {
return config;
}
private static SearchScript.LeafFactory createScript(Script script, QueryShardContext context) {
private static AggregationScript.LeafFactory createScript(Script script, QueryShardContext context) {
if (script == null) {
return null;
} else {
SearchScript.Factory factory = context.getScriptService().compile(script, SearchScript.AGGS_CONTEXT);
AggregationScript.Factory factory = context.getScriptService().compile(script, AggregationScript.CONTEXT);
return factory.newFactory(script.getParams(), context.lookup());
}
}
@ -135,7 +135,7 @@ public class ValuesSourceConfig<VS extends ValuesSource> {
private final ValuesSourceType valueSourceType;
private FieldContext fieldContext;
private SearchScript.LeafFactory script;
private AggregationScript.LeafFactory script;
private ValueType scriptValueType;
private boolean unmapped = false;
private DocValueFormat format = DocValueFormat.RAW;
@ -154,7 +154,7 @@ public class ValuesSourceConfig<VS extends ValuesSource> {
return fieldContext;
}
public SearchScript.LeafFactory script() {
public AggregationScript.LeafFactory script() {
return script;
}
@ -171,7 +171,7 @@ public class ValuesSourceConfig<VS extends ValuesSource> {
return this;
}
public ValuesSourceConfig<VS> script(SearchScript.LeafFactory script) {
public ValuesSourceConfig<VS> script(AggregationScript.LeafFactory script) {
this.script = script;
return this;
}

View File

@ -23,7 +23,7 @@ import org.elasticsearch.common.lucene.ScorerAware;
import org.elasticsearch.common.util.CollectionUtils;
import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
import org.elasticsearch.index.fielddata.SortingBinaryDocValues;
import org.elasticsearch.script.SearchScript;
import org.elasticsearch.script.AggregationScript;
import java.io.IOException;
import java.lang.reflect.Array;
@ -34,9 +34,9 @@ import java.util.Collection;
*/
public class ScriptBytesValues extends SortingBinaryDocValues implements ScorerAware {
private final SearchScript script;
private final AggregationScript script;
public ScriptBytesValues(SearchScript script) {
public ScriptBytesValues(AggregationScript script) {
super();
this.script = script;
}
@ -53,7 +53,7 @@ public class ScriptBytesValues extends SortingBinaryDocValues implements ScorerA
@Override
public boolean advanceExact(int doc) throws IOException {
script.setDocument(doc);
final Object value = script.run();
final Object value = script.execute();
if (value == null) {
return false;
} else if (value.getClass().isArray()) {

View File

@ -21,8 +21,8 @@ package org.elasticsearch.search.aggregations.support.values;
import org.apache.lucene.search.Scorable;
import org.elasticsearch.common.lucene.ScorerAware;
import org.elasticsearch.index.fielddata.SortingNumericDoubleValues;
import org.elasticsearch.script.AggregationScript;
import org.elasticsearch.script.JodaCompatibleZonedDateTime;
import org.elasticsearch.script.SearchScript;
import org.elasticsearch.search.aggregations.AggregationExecutionException;
import org.joda.time.ReadableInstant;
@ -36,9 +36,9 @@ import java.util.Collection;
*/
public class ScriptDoubleValues extends SortingNumericDoubleValues implements ScorerAware {
final SearchScript script;
final AggregationScript script;
public ScriptDoubleValues(SearchScript script) {
public ScriptDoubleValues(AggregationScript script) {
super();
this.script = script;
}
@ -46,7 +46,7 @@ public class ScriptDoubleValues extends SortingNumericDoubleValues implements Sc
@Override
public boolean advanceExact(int target) throws IOException {
script.setDocument(target);
final Object value = script.run();
final Object value = script.execute();
if (value == null) {
return false;

View File

@ -22,8 +22,8 @@ import org.apache.lucene.search.Scorable;
import org.apache.lucene.util.LongValues;
import org.elasticsearch.common.lucene.ScorerAware;
import org.elasticsearch.index.fielddata.AbstractSortingNumericDocValues;
import org.elasticsearch.script.AggregationScript;
import org.elasticsearch.script.JodaCompatibleZonedDateTime;
import org.elasticsearch.script.SearchScript;
import org.elasticsearch.search.aggregations.AggregationExecutionException;
import org.joda.time.ReadableInstant;
@ -38,9 +38,9 @@ import java.util.Iterator;
*/
public class ScriptLongValues extends AbstractSortingNumericDocValues implements ScorerAware {
final SearchScript script;
final AggregationScript script;
public ScriptLongValues(SearchScript script) {
public ScriptLongValues(AggregationScript script) {
super();
this.script = script;
}
@ -48,7 +48,7 @@ public class ScriptLongValues extends AbstractSortingNumericDocValues implements
@Override
public boolean advanceExact(int target) throws IOException {
script.setDocument(target);
final Object value = script.run();
final Object value = script.execute();
if (value == null) {
return false;

View File

@ -166,7 +166,7 @@ public class ScriptServiceTests extends ESTestCase {
buildScriptService(Settings.EMPTY);
assertCompileAccepted("painless", "script", ScriptType.INLINE, SearchScript.CONTEXT);
assertCompileAccepted("painless", "script", ScriptType.INLINE, SearchScript.AGGS_CONTEXT);
assertCompileAccepted("painless", "script", ScriptType.INLINE, AggregationScript.CONTEXT);
assertCompileAccepted("painless", "script", ScriptType.INLINE, UpdateScript.CONTEXT);
assertCompileAccepted("painless", "script", ScriptType.INLINE, IngestScript.CONTEXT);
}
@ -186,7 +186,7 @@ public class ScriptServiceTests extends ESTestCase {
buildScriptService(builder.build());
assertCompileAccepted("painless", "script", ScriptType.INLINE, SearchScript.CONTEXT);
assertCompileAccepted("painless", "script", ScriptType.INLINE, SearchScript.AGGS_CONTEXT);
assertCompileAccepted("painless", "script", ScriptType.INLINE, AggregationScript.CONTEXT);
assertCompileRejected("painless", "script", ScriptType.INLINE, UpdateScript.CONTEXT);
}
@ -205,7 +205,7 @@ public class ScriptServiceTests extends ESTestCase {
buildScriptService(builder.build());
assertCompileRejected("painless", "script", ScriptType.INLINE, SearchScript.CONTEXT);
assertCompileRejected("painless", "script", ScriptType.INLINE, SearchScript.AGGS_CONTEXT);
assertCompileRejected("painless", "script", ScriptType.INLINE, AggregationScript.CONTEXT);
}
public void testCompileNonRegisteredContext() throws IOException {

View File

@ -25,7 +25,6 @@ import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.index.fielddata.ScriptDocValues;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.plugins.Plugin;
import org.elasticsearch.script.ScoreAccessor;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.aggregations.AggregationTestScriptsPlugin;
@ -107,7 +106,7 @@ public class DoubleTermsIT extends AbstractTermsTestCase {
});
scripts.put("ceil(_score.doubleValue()/3)", vars -> {
ScoreAccessor score = (ScoreAccessor) vars.get("_score");
Number score = (Number) vars.get("_score");
return Math.ceil(score.doubleValue() / 3);
});

View File

@ -20,36 +20,48 @@
package org.elasticsearch.search.aggregations.support;
import com.carrotsearch.randomizedtesting.generators.RandomStrings;
import java.util.Collections;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.Scorable;
import org.apache.lucene.util.BytesRef;
import org.elasticsearch.script.SearchScript;
import org.elasticsearch.common.Strings;
import org.elasticsearch.script.AggregationScript;
import org.elasticsearch.search.aggregations.support.values.ScriptBytesValues;
import org.elasticsearch.search.aggregations.support.values.ScriptDoubleValues;
import org.elasticsearch.search.aggregations.support.values.ScriptLongValues;
import org.elasticsearch.search.lookup.LeafSearchLookup;
import org.elasticsearch.search.lookup.SearchLookup;
import org.elasticsearch.test.ESTestCase;
import java.io.IOException;
import java.util.Arrays;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
public class ScriptValuesTests extends ESTestCase {
private static class FakeSearchScript extends SearchScript {
private static class FakeAggregationScript extends AggregationScript {
private final Object[][] values;
int index;
FakeSearchScript(Object[][] values) {
super(null, null, null);
FakeAggregationScript(Object[][] values) {
super(Collections.emptyMap(), new SearchLookup(null, null, Strings.EMPTY_ARRAY) {
@Override
public LeafSearchLookup getLeafSearchLookup(LeafReaderContext context) {
LeafSearchLookup leafSearchLookup = mock(LeafSearchLookup.class);
when(leafSearchLookup.asMap()).thenReturn(Collections.emptyMap());
return leafSearchLookup;
}
}, null);
this.values = values;
index = -1;
}
@Override
public void setNextVar(String name, Object value) {
}
@Override
public Object run() {
public Object execute() {
// Script values are supposed to support null, single values, arrays and collections
final Object[] values = this.values[index];
if (values.length <= 1 && randomBoolean()) {
@ -89,7 +101,7 @@ public class ScriptValuesTests extends ESTestCase {
Arrays.sort(longs);
values[i] = longs;
}
FakeSearchScript script = new FakeSearchScript(values);
FakeAggregationScript script = new FakeAggregationScript(values);
ScriptLongValues scriptValues = new ScriptLongValues(script);
for (int i = 0; i < values.length; ++i) {
assertEquals(values[i].length > 0, scriptValues.advanceExact(i));
@ -112,7 +124,7 @@ public class ScriptValuesTests extends ESTestCase {
Arrays.sort(booleans);
values[i] = booleans;
}
FakeSearchScript script = new FakeSearchScript(values);
FakeAggregationScript script = new FakeAggregationScript(values);
ScriptLongValues scriptValues = new ScriptLongValues(script);
for (int i = 0; i < values.length; ++i) {
assertEquals(values[i].length > 0, scriptValues.advanceExact(i));
@ -135,7 +147,7 @@ public class ScriptValuesTests extends ESTestCase {
Arrays.sort(doubles);
values[i] = doubles;
}
FakeSearchScript script = new FakeSearchScript(values);
FakeAggregationScript script = new FakeAggregationScript(values);
ScriptDoubleValues scriptValues = new ScriptDoubleValues(script);
for (int i = 0; i < values.length; ++i) {
assertEquals(values[i].length > 0, scriptValues.advanceExact(i));
@ -158,7 +170,7 @@ public class ScriptValuesTests extends ESTestCase {
Arrays.sort(strings);
values[i] = strings;
}
FakeSearchScript script = new FakeSearchScript(values);
FakeAggregationScript script = new FakeAggregationScript(values);
ScriptBytesValues scriptValues = new ScriptBytesValues(script);
for (int i = 0; i < values.length; ++i) {
assertEquals(values[i].length > 0, scriptValues.advanceExact(i));

View File

@ -28,7 +28,6 @@ import org.elasticsearch.index.query.MatchAllQueryBuilder;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder.FilterFunctionBuilder;
import org.elasticsearch.plugins.Plugin;
import org.elasticsearch.script.MockScriptPlugin;
import org.elasticsearch.script.ScoreAccessor;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
@ -73,7 +72,7 @@ public class FunctionScoreIT extends ESIntegTestCase {
protected Map<String, Function<Map<String, Object>, Object>> pluginScripts() {
Map<String, Function<Map<String, Object>, Object>> scripts = new HashMap<>();
scripts.put("1", vars -> 1.0d);
scripts.put("get score value", vars -> ((ScoreAccessor) vars.get("_score")).doubleValue());
scripts.put("get score value", vars -> ((Number) vars.get("_score")).doubleValue());
scripts.put("return (doc['num'].value)", vars -> {
Map<?, ?> doc = (Map) vars.get("doc");
ScriptDocValues.Longs num = (ScriptDocValues.Longs) doc.get("num");

View File

@ -105,6 +105,29 @@ public class MockScriptEngine implements ScriptEngine {
}
};
return context.factoryClazz.cast(factory);
} else if(context.instanceClazz.equals(AggregationScript.class)) {
AggregationScript.Factory factory = (parameters, lookup) -> new AggregationScript.LeafFactory() {
@Override
public AggregationScript newInstance(final LeafReaderContext ctx) {
return new AggregationScript(parameters, lookup, ctx) {
@Override
public Object execute() {
Map<String, Object> vars = new HashMap<>(parameters);
vars.put("params", parameters);
vars.put("doc", getDoc());
vars.put("_score", get_score());
vars.put("_value", get_value());
return script.apply(vars);
}
};
}
@Override
public boolean needs_score() {
return true;
}
};
return context.factoryClazz.cast(factory);
} else if (context.instanceClazz.equals(IngestScript.class)) {
IngestScript.Factory factory = vars ->
new IngestScript(vars) {
@ -307,11 +330,6 @@ public class MockScriptEngine implements ScriptEngine {
return script.apply(ctx);
}
@Override
public long runAsLong() {
return ((Number) run()).longValue();
}
@Override
public double runAsDouble() {
return ((Number) run()).doubleValue();

View File

@ -8,6 +8,7 @@ package org.elasticsearch.xpack.sql.plugin;
import org.elasticsearch.painless.spi.PainlessExtension;
import org.elasticsearch.painless.spi.Whitelist;
import org.elasticsearch.painless.spi.WhitelistLoader;
import org.elasticsearch.script.AggregationScript;
import org.elasticsearch.script.BucketAggregationSelectorScript;
import org.elasticsearch.script.FilterScript;
import org.elasticsearch.script.ScriptContext;
@ -28,7 +29,7 @@ public class SqlPainlessExtension implements PainlessExtension {
Map<ScriptContext<?>, List<Whitelist>> whitelist = new HashMap<>();
List<Whitelist> list = singletonList(WHITELIST);
whitelist.put(FilterScript.CONTEXT, list);
whitelist.put(SearchScript.AGGS_CONTEXT, list);
whitelist.put(AggregationScript.CONTEXT, list);
whitelist.put(SearchScript.CONTEXT, list);
whitelist.put(SearchScript.SCRIPT_SORT_CONTEXT, list);
whitelist.put(BucketAggregationSelectorScript.CONTEXT, list);