diff --git a/modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionNumberSortScript.java b/modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionNumberSortScript.java new file mode 100644 index 0000000000000..4a7bbc6182e2a --- /dev/null +++ b/modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionNumberSortScript.java @@ -0,0 +1,91 @@ +/* + * 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.GeneralScriptException; +import org.elasticsearch.script.NumberSortScript; + +/** + * A bridge to evaluate an {@link Expression} against {@link Bindings} in the context + * of a {@link NumberSortScript}. + */ +class ExpressionNumberSortScript implements NumberSortScript.LeafFactory { + + final Expression exprScript; + final SimpleBindings bindings; + final DoubleValuesSource source; + final boolean needsScores; + + ExpressionNumberSortScript(Expression e, SimpleBindings b, boolean needsScores) { + exprScript = e; + bindings = b; + source = exprScript.getDoubleValuesSource(bindings); + this.needsScores = needsScores; + } + + @Override + public NumberSortScript newInstance(final LeafReaderContext leaf) throws IOException { + return new NumberSortScript() { + // Fake the scorer until setScorer is called. + DoubleValues values = source.getValues(leaf, new DoubleValues() { + @Override + public double doubleValue() { + return 0.0D; + } + + @Override + public boolean advanceExact(int doc) { + return true; + } + }); + + @Override + public double 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 boolean needs_score() { + return needsScores; + } + +} diff --git a/modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java b/modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java index 14d701fad8b07..5d47c87a445a2 100644 --- a/modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java +++ b/modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java @@ -42,6 +42,7 @@ import org.elasticsearch.script.BucketAggregationSelectorScript; import org.elasticsearch.script.ClassPermission; import org.elasticsearch.script.FilterScript; +import org.elasticsearch.script.NumberSortScript; import org.elasticsearch.script.ScoreScript; import org.elasticsearch.script.ScriptContext; import org.elasticsearch.script.ScriptEngine; @@ -135,6 +136,9 @@ public boolean execute() { } else if (context.instanceClazz.equals(AggregationScript.class)) { AggregationScript.Factory factory = (p, lookup) -> newAggregationScript(expr, lookup, p); return context.factoryClazz.cast(factory); + } else if (context.instanceClazz.equals(NumberSortScript.class)) { + NumberSortScript.Factory factory = (p, lookup) -> newSortScript(expr, lookup, p); + return context.factoryClazz.cast(factory); } throw new IllegalArgumentException("expression engine does not know how to handle script context [" + context.name + "]"); } @@ -187,7 +191,6 @@ private SearchScript.LeafFactory newSearchScript(Expression expr, SearchLookup l // 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 { @@ -205,6 +208,33 @@ private SearchScript.LeafFactory newSearchScript(Expression expr, SearchLookup l return new ExpressionSearchScript(expr, bindings, specialValue, needsScores); } + private NumberSortScript.LeafFactory newSortScript(Expression expr, SearchLookup lookup, @Nullable Map 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(); + boolean needsScores = false; + for (String variable : expr.variables) { + try { + if (variable.equals("_score")) { + bindings.add(new SortField("_score", SortField.Type.SCORE)); + needsScores = true; + } 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); + needsScores |= valueSource.getSortField(false).needsScores(); + 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 ExpressionNumberSortScript(expr, bindings, needsScores); + } + private TermsSetQueryScript.LeafFactory newTermsSetQueryScript(Expression expr, SearchLookup lookup, @Nullable Map vars) { // NOTE: if we need to do anything complicated with bindings in the future, we can just extend Bindings, diff --git a/modules/lang-expression/src/test/java/org/elasticsearch/script/expression/ExpressionNumberSortScriptTests.java b/modules/lang-expression/src/test/java/org/elasticsearch/script/expression/ExpressionNumberSortScriptTests.java new file mode 100644 index 0000000000000..301fd2d4db70c --- /dev/null +++ b/modules/lang-expression/src/test/java/org/elasticsearch/script/expression/ExpressionNumberSortScriptTests.java @@ -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.NumberSortScript; +import org.elasticsearch.script.ScriptException; +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 ExpressionNumberSortScriptTests 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 NumberSortScript.LeafFactory compile(String expression) { + NumberSortScript.Factory factory = + service.compile(null, expression, NumberSortScript.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 { + NumberSortScript script = compile("doc['field'].value").newInstance(null); + script.setDocument(1); + + double result = script.execute(); + assertEquals(2.718, result, 0.0); + } + + public void testFieldAccessWithFieldAlias() throws IOException { + NumberSortScript script = compile("doc['alias'].value").newInstance(null); + script.setDocument(1); + + double result = script.execute(); + assertEquals(2.718, result, 0.0); + } +} diff --git a/modules/lang-painless/src/test/java/org/elasticsearch/painless/NeedsScoreTests.java b/modules/lang-painless/src/test/java/org/elasticsearch/painless/NeedsScoreTests.java index 86f2af32d1632..eeb636d6697c6 100644 --- a/modules/lang-painless/src/test/java/org/elasticsearch/painless/NeedsScoreTests.java +++ b/modules/lang-painless/src/test/java/org/elasticsearch/painless/NeedsScoreTests.java @@ -23,8 +23,8 @@ import org.elasticsearch.index.IndexService; import org.elasticsearch.index.query.QueryShardContext; import org.elasticsearch.painless.spi.Whitelist; +import org.elasticsearch.script.NumberSortScript; import org.elasticsearch.script.ScriptContext; -import org.elasticsearch.script.SearchScript; import org.elasticsearch.search.lookup.SearchLookup; import org.elasticsearch.test.ESSingleNodeTestCase; @@ -43,25 +43,25 @@ public void testNeedsScores() { IndexService index = createIndex("test", Settings.EMPTY, "type", "d", "type=double"); Map, List> contexts = new HashMap<>(); - contexts.put(SearchScript.CONTEXT, Whitelist.BASE_WHITELISTS); + contexts.put(NumberSortScript.CONTEXT, Whitelist.BASE_WHITELISTS); PainlessScriptEngine service = new PainlessScriptEngine(Settings.EMPTY, contexts); QueryShardContext shardContext = index.newQueryShardContext(0, null, () -> 0, null); SearchLookup lookup = new SearchLookup(index.mapperService(), shardContext::getForField, null); - SearchScript.Factory factory = service.compile(null, "1.2", SearchScript.CONTEXT, Collections.emptyMap()); - SearchScript.LeafFactory ss = factory.newFactory(Collections.emptyMap(), lookup); + NumberSortScript.Factory factory = service.compile(null, "1.2", NumberSortScript.CONTEXT, Collections.emptyMap()); + NumberSortScript.LeafFactory ss = factory.newFactory(Collections.emptyMap(), lookup); assertFalse(ss.needs_score()); - factory = service.compile(null, "doc['d'].value", SearchScript.CONTEXT, Collections.emptyMap()); + factory = service.compile(null, "doc['d'].value", NumberSortScript.CONTEXT, Collections.emptyMap()); ss = factory.newFactory(Collections.emptyMap(), lookup); assertFalse(ss.needs_score()); - factory = service.compile(null, "1/_score", SearchScript.CONTEXT, Collections.emptyMap()); + factory = service.compile(null, "1/_score", NumberSortScript.CONTEXT, Collections.emptyMap()); ss = factory.newFactory(Collections.emptyMap(), lookup); assertTrue(ss.needs_score()); - factory = service.compile(null, "doc['d'].value * _score", SearchScript.CONTEXT, Collections.emptyMap()); + factory = service.compile(null, "doc['d'].value * _score", NumberSortScript.CONTEXT, Collections.emptyMap()); ss = factory.newFactory(Collections.emptyMap(), lookup); assertTrue(ss.needs_score()); } diff --git a/server/src/main/java/org/elasticsearch/script/AbstractSortScript.java b/server/src/main/java/org/elasticsearch/script/AbstractSortScript.java new file mode 100644 index 0000000000000..1d8de9f95f40a --- /dev/null +++ b/server/src/main/java/org/elasticsearch/script/AbstractSortScript.java @@ -0,0 +1,111 @@ +/* + * 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; + +abstract class AbstractSortScript implements ScorerAware { + + private static final Map DEPRECATIONS; + + static { + Map deprecations = new HashMap<>(); + deprecations.put( + "doc", + "Accessing variable [doc] via [params.doc] from within a sort-script " + + "is deprecated in favor of directly accessing [doc]." + ); + deprecations.put( + "_doc", + "Accessing variable [doc] via [params._doc] from within a sort-script " + + "is deprecated in favor of directly accessing [doc]." + ); + DEPRECATIONS = Collections.unmodifiableMap(deprecations); + } + + /** + * The generic runtime parameters for the script. + */ + private final Map params; + + /** A scorer that will return the score for the current document when the script is run. */ + private Scorable scorer; + + /** + * A leaf lookup for the bound segment this script will operate on. + */ + private final LeafSearchLookup leafLookup; + + AbstractSortScript(Map params, SearchLookup lookup, LeafReaderContext leafContext) { + this.leafLookup = lookup.getLeafSearchLookup(leafContext); + Map parameters = new HashMap<>(params); + parameters.putAll(leafLookup.asMap()); + this.params = new ParameterMap(parameters, DEPRECATIONS); + } + + protected AbstractSortScript() { + this.params = null; + this.leafLookup = null; + } + + /** + * Return the parameters for this script. + */ + public Map getParams() { + return params; + } + + @Override + public void setScorer(Scorable scorer) { + this.scorer = scorer; + } + + /** Return the score of the current document. */ + public double get_score() { + try { + return scorer.score(); + } catch (IOException e) { + throw new ElasticsearchException("couldn't lookup score", e); + } + } + + /** + * The doc lookup for the Lucene segment this script was created for. + */ + public Map> getDoc() { + return leafLookup.doc(); + } + + /** + * Set the current document to run the script on next. + */ + public void setDocument(int docid) { + leafLookup.setDocument(docid); + } +} diff --git a/server/src/main/java/org/elasticsearch/script/NumberSortScript.java b/server/src/main/java/org/elasticsearch/script/NumberSortScript.java new file mode 100644 index 0000000000000..d0b3fdbed363e --- /dev/null +++ b/server/src/main/java/org/elasticsearch/script/NumberSortScript.java @@ -0,0 +1,60 @@ +/* + * 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.Map; +import org.apache.lucene.index.LeafReaderContext; +import org.elasticsearch.search.lookup.SearchLookup; + +public abstract class NumberSortScript extends AbstractSortScript { + + public static final String[] PARAMETERS = {}; + + public static final ScriptContext CONTEXT = new ScriptContext<>("number_sort", Factory.class); + + public NumberSortScript(Map params, SearchLookup lookup, LeafReaderContext leafContext) { + super(params, lookup, leafContext); + } + + protected NumberSortScript() { + super(); + } + + public abstract double execute(); + + /** + * A factory to construct {@link NumberSortScript} instances. + */ + public interface LeafFactory { + NumberSortScript 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 NumberSortScript} factories for a specific index. + */ + public interface Factory { + LeafFactory newFactory(Map params, SearchLookup lookup); + } +} diff --git a/server/src/main/java/org/elasticsearch/script/ScriptModule.java b/server/src/main/java/org/elasticsearch/script/ScriptModule.java index ddc090251a392..24dd491e3a183 100644 --- a/server/src/main/java/org/elasticsearch/script/ScriptModule.java +++ b/server/src/main/java/org/elasticsearch/script/ScriptModule.java @@ -43,7 +43,8 @@ public class ScriptModule { SearchScript.CONTEXT, AggregationScript.CONTEXT, ScoreScript.CONTEXT, - SearchScript.SCRIPT_SORT_CONTEXT, + NumberSortScript.CONTEXT, + StringSortScript.CONTEXT, TermsSetQueryScript.CONTEXT, UpdateScript.CONTEXT, BucketAggregationScript.CONTEXT, diff --git a/server/src/main/java/org/elasticsearch/script/SearchScript.java b/server/src/main/java/org/elasticsearch/script/SearchScript.java index 496af2fbdd5d6..2fd439564a61f 100644 --- a/server/src/main/java/org/elasticsearch/script/SearchScript.java +++ b/server/src/main/java/org/elasticsearch/script/SearchScript.java @@ -139,7 +139,4 @@ public interface Factory { /** The context used to compile {@link SearchScript} factories. */ public static final ScriptContext CONTEXT = new ScriptContext<>("search", Factory.class); - // TODO: remove these contexts when it has its own interface - // Can return a double. (For ScriptSortType#NUMBER only, for ScriptSortType#STRING normal CONTEXT should be used) - public static final ScriptContext SCRIPT_SORT_CONTEXT = new ScriptContext<>("sort", Factory.class); } diff --git a/server/src/main/java/org/elasticsearch/script/StringSortScript.java b/server/src/main/java/org/elasticsearch/script/StringSortScript.java new file mode 100644 index 0000000000000..1c6c47dd21552 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/script/StringSortScript.java @@ -0,0 +1,51 @@ +/* + * 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.Map; +import org.apache.lucene.index.LeafReaderContext; +import org.elasticsearch.search.lookup.SearchLookup; + +public abstract class StringSortScript extends AbstractSortScript { + + public static final String[] PARAMETERS = {}; + + public static final ScriptContext CONTEXT = new ScriptContext<>("string_sort", Factory.class); + + public StringSortScript(Map params, SearchLookup lookup, LeafReaderContext leafContext) { + super(params, lookup, leafContext); + } + + public abstract String execute(); + + /** + * A factory to construct {@link StringSortScript} instances. + */ + public interface LeafFactory { + StringSortScript newInstance(LeafReaderContext ctx) throws IOException; + } + + /** + * A factory to construct stateful {@link StringSortScript} factories for a specific index. + */ + public interface Factory { + LeafFactory newFactory(Map params, SearchLookup lookup); + } +} diff --git a/server/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java b/server/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java index e425755a55edd..95478e083243a 100644 --- a/server/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java @@ -32,7 +32,6 @@ import org.elasticsearch.common.io.stream.Writeable; import org.elasticsearch.common.logging.DeprecationLogger; import org.elasticsearch.common.logging.Loggers; -import org.elasticsearch.common.util.CollectionUtils; import org.elasticsearch.common.xcontent.ConstructingObjectParser; import org.elasticsearch.common.xcontent.ObjectParser.ValueType; import org.elasticsearch.common.xcontent.XContentBuilder; @@ -51,7 +50,8 @@ import org.elasticsearch.index.query.QueryShardContext; import org.elasticsearch.index.query.QueryShardException; import org.elasticsearch.script.Script; -import org.elasticsearch.script.SearchScript; +import org.elasticsearch.script.NumberSortScript; +import org.elasticsearch.script.StringSortScript; import org.elasticsearch.search.DocValueFormat; import org.elasticsearch.search.MultiValueMode; @@ -305,9 +305,6 @@ public static ScriptSortBuilder fromXContent(XContentParser parser, String eleme @Override public SortFieldAndFormat build(QueryShardContext context) throws IOException { - final SearchScript.Factory factory = context.getScriptService().compile(script, SearchScript.SCRIPT_SORT_CONTEXT); - final SearchScript.LeafFactory searchScript = factory.newFactory(script.getParams(), context.lookup()); - MultiValueMode valueMode = null; if (sortMode != null) { valueMode = MultiValueMode.fromString(sortMode.toString()); @@ -336,8 +333,10 @@ public SortFieldAndFormat build(QueryShardContext context) throws IOException { final IndexFieldData.XFieldComparatorSource fieldComparatorSource; switch (type) { case STRING: + final StringSortScript.Factory factory = context.getScriptService().compile(script, StringSortScript.CONTEXT); + final StringSortScript.LeafFactory searchScript = factory.newFactory(script.getParams(), context.lookup()); fieldComparatorSource = new BytesRefFieldComparatorSource(null, null, valueMode, nested) { - SearchScript leafScript; + StringSortScript leafScript; @Override protected SortedBinaryDocValues getValues(LeafReaderContext context) throws IOException { leafScript = searchScript.newInstance(context); @@ -350,9 +349,7 @@ public boolean advanceExact(int doc) throws IOException { } @Override public BytesRef binaryValue() { - final Object run = leafScript.run(); - CollectionUtils.ensureNoSelfReferences(run, "ScriptSortBuilder leaf script"); - spare.copyChars(run.toString()); + spare.copyChars(leafScript.execute()); return spare.get(); } }; @@ -365,11 +362,13 @@ protected void setScorer(Scorable scorer) { }; break; case NUMBER: + final NumberSortScript.Factory numberSortFactory = context.getScriptService().compile(script, NumberSortScript.CONTEXT); + final NumberSortScript.LeafFactory numberSortScript = numberSortFactory.newFactory(script.getParams(), context.lookup()); fieldComparatorSource = new DoubleValuesComparatorSource(null, Double.MAX_VALUE, valueMode, nested) { - SearchScript leafScript; + NumberSortScript leafScript; @Override protected SortedNumericDoubleValues getValues(LeafReaderContext context) throws IOException { - leafScript = searchScript.newInstance(context); + leafScript = numberSortScript.newInstance(context); final NumericDoubleValues values = new NumericDoubleValues() { @Override public boolean advanceExact(int doc) throws IOException { @@ -378,7 +377,7 @@ public boolean advanceExact(int doc) throws IOException { } @Override public double doubleValue() { - return leafScript.runAsDouble(); + return leafScript.execute(); } }; return FieldData.singleton(values); diff --git a/test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java b/test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java index 9e89bf2b59d91..64a9d97e3f429 100644 --- a/test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java +++ b/test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java @@ -105,6 +105,47 @@ public Number execute() { } }; return context.factoryClazz.cast(factory); + } else if (context.instanceClazz.equals(NumberSortScript.class)) { + NumberSortScript.Factory factory = (parameters, lookup) -> new NumberSortScript.LeafFactory() { + @Override + public NumberSortScript newInstance(final LeafReaderContext ctx) { + return new NumberSortScript(parameters, lookup, ctx) { + @Override + public double execute() { + Map vars = new HashMap<>(parameters); + vars.put("params", parameters); + vars.put("doc", getDoc()); + return ((Number) script.apply(vars)).doubleValue(); + } + }; + } + + @Override + public boolean needs_score() { + return false; + } + }; + return context.factoryClazz.cast(factory); + } else if (context.instanceClazz.equals(StringSortScript.class)) { + StringSortScript.Factory factory = (parameters, lookup) -> (StringSortScript.LeafFactory) ctx + -> new StringSortScript(parameters, lookup, ctx) { + @Override + public String execute() { + Map vars = new HashMap<>(parameters); + vars.put("params", parameters); + vars.put("doc", getDoc()); + return String.valueOf(script.apply(vars)); + } + }; + return context.factoryClazz.cast(factory); + } else if (context.instanceClazz.equals(IngestScript.class)) { + IngestScript.Factory factory = vars -> new IngestScript(vars) { + @Override + public void execute(Map ctx) { + script.apply(ctx); + } + }; + return context.factoryClazz.cast(factory); } else if(context.instanceClazz.equals(AggregationScript.class)) { AggregationScript.Factory factory = (parameters, lookup) -> new AggregationScript.LeafFactory() { @Override diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/SqlPainlessExtension.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/SqlPainlessExtension.java index 4a688eb333423..f0ca0cb87308d 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/SqlPainlessExtension.java +++ b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/SqlPainlessExtension.java @@ -17,6 +17,8 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import org.elasticsearch.script.NumberSortScript; +import org.elasticsearch.script.StringSortScript; import static java.util.Collections.singletonList; @@ -31,7 +33,8 @@ public Map, List> getContextWhitelists() { whitelist.put(FilterScript.CONTEXT, list); whitelist.put(AggregationScript.CONTEXT, list); whitelist.put(SearchScript.CONTEXT, list); - whitelist.put(SearchScript.SCRIPT_SORT_CONTEXT, list); + whitelist.put(NumberSortScript.CONTEXT, list); + whitelist.put(StringSortScript.CONTEXT, list); whitelist.put(BucketAggregationSelectorScript.CONTEXT, list); return whitelist; }