Skip to content

Commit

Permalink
SCRIPTING: Move Aggregation Script Context to its own class (#33820)
Browse files Browse the repository at this point in the history
* SCRIPTING: Move Aggregation Script Context to its own class
  • Loading branch information
original-brownbear authored Oct 15, 2018
1 parent 5115262 commit ebca273
Show file tree
Hide file tree
Showing 18 changed files with 384 additions and 76 deletions.
Original file line number Diff line number Diff line change
@@ -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;
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
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;
Expand Down Expand Up @@ -131,6 +132,9 @@ public boolean execute() {
} 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 + "]");
}
Expand Down Expand Up @@ -224,6 +228,37 @@ private TermsSetQueryScript.LeafFactory newTermsSetQueryScript(Expression expr,
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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -75,9 +75,6 @@ public boolean advanceExact(int doc) throws IOException {
@Override
public Object run() { return Double.valueOf(runAsDouble()); }

@Override
public long runAsLong() { return (long)runAsDouble(); }

@Override
public double runAsDouble() {
try {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -109,8 +109,4 @@ public double runAsDouble() {
return ((Number)run()).doubleValue();
}

@Override
public long runAsLong() {
return ((Number)run()).longValue();
}
}
165 changes: 165 additions & 0 deletions server/src/main/java/org/elasticsearch/script/AggregationScript.java
Original file line number Diff line number Diff line change
@@ -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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
* <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 {
Expand Down Expand Up @@ -114,10 +114,6 @@ public void setNextAggregationValue(Object value) {

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();
Expand All @@ -144,7 +140,6 @@ public interface Factory {
/** 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);
}
Loading

0 comments on commit ebca273

Please sign in to comment.