Skip to content

Commit

Permalink
SCRIPTING: Add Expr. Compile for TermSetQuery Ctx.
Browse files Browse the repository at this point in the history
* Follow up to elastic#33602 adding the ability to compile TermsSetQuery
scripts with the expressions engine in the same way we support
SearchScript in Expressions
   * Duplicated the code here for now to make the change less complex,
 the only difference to SearchScript is that `_score` and `_value` are not handled for TermsSetQuery
  • Loading branch information
original-brownbear committed Sep 19, 2018
1 parent d9947c6 commit 20a66f9
Show file tree
Hide file tree
Showing 4 changed files with 308 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
import org.elasticsearch.script.ScriptEngine;
import org.elasticsearch.script.ScriptException;
import org.elasticsearch.script.SearchScript;
import org.elasticsearch.script.TermsSetQueryScript;
import org.elasticsearch.search.lookup.SearchLookup;

import java.io.IOException;
Expand Down Expand Up @@ -132,6 +133,9 @@ public boolean execute() {
} else if (context.instanceClazz.equals(ScoreScript.class)) {
ScoreScript.Factory factory = (p, lookup) -> newScoreScript(expr, lookup, p);
return context.factoryClazz.cast(factory);
} else if (context.instanceClazz.equals(TermsSetQueryScript.class)) {
TermsSetQueryScript.Factory factory = (p, lookup) -> newTermsSetQueryScript(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 @@ -292,6 +296,118 @@ private SearchScript.LeafFactory newSearchScript(Expression expr, SearchLookup l
return new ExpressionSearchScript(expr, bindings, specialValue, needsScores);
}

private TermsSetQueryScript.LeafFactory newTermsSetQueryScript(Expression expr, SearchLookup lookup, @Nullable Map<String, Object> vars) {
MapperService mapper = lookup.doc().mapperService();
// 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();
for (String variable : expr.variables) {
try {
if (vars != null && vars.containsKey(variable)) {
// NOTE: by checking for the variable in vars first, it allows masking document fields with a global constant,
// but if we were to reverse it, we could provide a way to supply dynamic defaults for documents missing the field?
Object value = vars.get(variable);
if (value instanceof Number) {
bindings.add(variable, new DoubleConstValueSource(((Number) value).doubleValue()).asDoubleValuesSource());
} else {
throw new ParseException("Parameter [" + variable + "] must be a numeric type", 0);
}

} else {
String fieldname = null;
String methodname = null;
String variablename = "value"; // .value is the default for doc['field'], its optional.
boolean dateAccessor = false; // true if the variable is of type doc['field'].date.xxx
VariableContext[] parts = VariableContext.parse(variable);
if (parts[0].text.equals("doc") == false) {
throw new ParseException("Unknown variable [" + parts[0].text + "]", 0);
}
if (parts.length < 2 || parts[1].type != VariableContext.Type.STR_INDEX) {
throw new ParseException("Variable 'doc' must be used with a specific field like: doc['myfield']", 3);
} else {
fieldname = parts[1].text;
}
if (parts.length == 3) {
if (parts[2].type == VariableContext.Type.METHOD) {
methodname = parts[2].text;
} else if (parts[2].type == VariableContext.Type.MEMBER) {
variablename = parts[2].text;
} else {
throw new IllegalArgumentException("Only member variables or member methods may be accessed on a field when not accessing the field directly");
}
}
if (parts.length > 3) {
// access to the .date "object" within the field
if (parts.length == 4 && ("date".equals(parts[2].text) || "getDate".equals(parts[2].text))) {
if (parts[3].type == VariableContext.Type.METHOD) {
methodname = parts[3].text;
dateAccessor = true;
} else if (parts[3].type == VariableContext.Type.MEMBER) {
variablename = parts[3].text;
dateAccessor = true;
}
}
if (!dateAccessor) {
throw new IllegalArgumentException("Variable [" + variable + "] does not follow an allowed format of either doc['field'] or doc['field'].method()");
}
}

MappedFieldType fieldType = mapper.fullName(fieldname);

if (fieldType == null) {
throw new ParseException("Field [" + fieldname + "] does not exist in mappings", 5);
}

IndexFieldData<?> fieldData = lookup.doc().getForField(fieldType);

// 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;
if (fieldType instanceof GeoPointFieldType) {
// geo
if (methodname == null) {
valueSource = GeoField.getVariable(fieldData, fieldname, variablename);
} else {
valueSource = GeoField.getMethod(fieldData, fieldname, methodname);
}
} else if (fieldType instanceof DateFieldMapper.DateFieldType) {
if (dateAccessor) {
// date object
if (methodname == null) {
valueSource = DateObject.getVariable(fieldData, fieldname, variablename);
} else {
valueSource = DateObject.getMethod(fieldData, fieldname, methodname);
}
} else {
// date field itself
if (methodname == null) {
valueSource = DateField.getVariable(fieldData, fieldname, variablename);
} else {
valueSource = DateField.getMethod(fieldData, fieldname, methodname);
}
}
} else if (fieldData instanceof IndexNumericFieldData) {
// number
if (methodname == null) {
valueSource = NumericField.getVariable(fieldData, fieldname, variablename);
} else {
valueSource = NumericField.getMethod(fieldData, fieldname, methodname);
}
} else {
throw new ParseException("Field [" + fieldname + "] must be numeric, date, or geopoint", 5);
}
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);
}
}
ReplaceableConstDoubleValueSource specialValue = null;
return new ExpressionTermSetQueryScript(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
@@ -0,0 +1,77 @@
/*
* 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 java.util.Collections;
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.GeneralScriptException;
import org.elasticsearch.script.TermsSetQueryScript;

/**
* A bridge to evaluate an {@link Expression} against {@link Bindings} in the context
* of a {@link TermsSetQueryScript}.
*/
class ExpressionTermSetQueryScript implements TermsSetQueryScript.LeafFactory {

final Expression exprScript;
final SimpleBindings bindings;
final DoubleValuesSource source;
final ReplaceableConstDoubleValueSource specialValue; // _value

ExpressionTermSetQueryScript(Expression e, SimpleBindings b, ReplaceableConstDoubleValueSource v) {
exprScript = e;
bindings = b;
source = exprScript.getDoubleValuesSource(bindings);
specialValue = v;
}

@Override
public TermsSetQueryScript newInstance(final LeafReaderContext leaf) throws IOException {
return new TermsSetQueryScript(Collections.emptyMap(), null, null) {
// Fake the scorer until setScorer is called.
DoubleValues values = source.getValues(leaf, null);

@Override
public Number 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);
}
}
};
}

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
/*
* 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 java.text.ParseException;
import java.util.Collections;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.index.fielddata.AtomicNumericFieldData;
import org.elasticsearch.index.fielddata.IndexNumericFieldData;
import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
import org.elasticsearch.index.mapper.MapperService;
import org.elasticsearch.index.mapper.NumberFieldMapper.NumberFieldType;
import org.elasticsearch.index.mapper.NumberFieldMapper.NumberType;
import org.elasticsearch.script.ScriptException;
import org.elasticsearch.script.TermsSetQueryScript;
import org.elasticsearch.search.lookup.SearchLookup;
import org.elasticsearch.test.ESTestCase;

import static org.mockito.Matchers.anyInt;
import static org.mockito.Matchers.anyObject;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

public class ExpressionTermsSetQueryTests extends ESTestCase {
private ExpressionScriptEngine service;
private SearchLookup lookup;

@Override
public void setUp() throws Exception {
super.setUp();

NumberFieldType fieldType = new NumberFieldType(NumberType.DOUBLE);
MapperService mapperService = mock(MapperService.class);
when(mapperService.fullName("field")).thenReturn(fieldType);
when(mapperService.fullName("alias")).thenReturn(fieldType);

SortedNumericDoubleValues doubleValues = mock(SortedNumericDoubleValues.class);
when(doubleValues.advanceExact(anyInt())).thenReturn(true);
when(doubleValues.nextValue()).thenReturn(2.718);

AtomicNumericFieldData atomicFieldData = mock(AtomicNumericFieldData.class);
when(atomicFieldData.getDoubleValues()).thenReturn(doubleValues);

IndexNumericFieldData fieldData = mock(IndexNumericFieldData.class);
when(fieldData.getFieldName()).thenReturn("field");
when(fieldData.load(anyObject())).thenReturn(atomicFieldData);

service = new ExpressionScriptEngine(Settings.EMPTY);
lookup = new SearchLookup(mapperService, ignored -> fieldData, null);
}

private TermsSetQueryScript.LeafFactory compile(String expression) {
TermsSetQueryScript.Factory factory =
service.compile(null, expression, TermsSetQueryScript.CONTEXT, Collections.emptyMap());
return factory.newFactory(Collections.emptyMap(), lookup);
}

public void testCompileError() {
ScriptException e = expectThrows(ScriptException.class, () -> {
compile("doc['field'].value * *@#)(@$*@#$ + 4");
});
assertTrue(e.getCause() instanceof ParseException);
}

public void testLinkError() {
ScriptException e = expectThrows(ScriptException.class, () -> {
compile("doc['nonexistent'].value * 5");
});
assertTrue(e.getCause() instanceof ParseException);
}

public void testFieldAccess() throws IOException {
TermsSetQueryScript script = compile("doc['field'].value").newInstance(null);
script.setDocument(1);

double result = script.execute().doubleValue();
assertEquals(2.718, result, 0.0);
}

public void testFieldAccessWithFieldAlias() throws IOException {
TermsSetQueryScript script = compile("doc['alias'].value").newInstance(null);
script.setDocument(1);

double result = script.execute().doubleValue();
assertEquals(2.718, result, 0.0);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -61,15 +61,19 @@ public abstract class TermsSetQueryScript {
private final LeafSearchLookup leafLookup;

public TermsSetQueryScript(Map<String, Object> params, SearchLookup lookup, LeafReaderContext leafContext) {
this.params = new ParameterMap(params, DEPRECATIONS);
this.leafLookup = lookup.getLeafSearchLookup(leafContext);
this.params = new ParameterMap(new HashMap<>(params), DEPRECATIONS);
if (leafContext == null) {
this.leafLookup = null;
} else {
this.leafLookup = lookup.getLeafSearchLookup(leafContext);
this.params.putAll(leafLookup.asMap());
}
}

/**
* Return the parameters for this script.
*/
public Map<String, Object> getParams() {
this.params.putAll(leafLookup.asMap());
return params;
}

Expand All @@ -84,7 +88,9 @@ public Map<String, ScriptDocValues<?>> getDoc() {
* Set the current document to run the script on next.
*/
public void setDocument(int docid) {
leafLookup.setDocument(docid);
if (leafLookup != null) {
leafLookup.setDocument(docid);
}
}

/**
Expand Down

0 comments on commit 20a66f9

Please sign in to comment.