Merge pull request #12695 from jpountz/enhancement/script_needs_scores

Allow scripts to expose whether they use the `_score`.

core/src/main/java/org/elasticsearch/common/lucene/search/function/ScriptScoreFunction.java

@@ -128,9 +128,7 @@ public class ScriptScoreFunction extends ScoreFunction {
 
     @Override
     public boolean needsScores() {
-        // Scripts might use _score so we return true here
+        return script.needsScores();
-        // TODO: Make scripts able to tell us whether they use scores
-        return true;
     }
 
     @Override

core/src/main/java/org/elasticsearch/script/NativeScriptEngineService.java

@@ -86,6 +86,10 @@ public class NativeScriptEngineService extends AbstractComponent implements Scri
                 script.setLookup(lookup.getLeafSearchLookup(context));
                 return script;
             }
+            @Override
+            public boolean needsScores() {
+                return scriptFactory.needsScores();
+            }
         };
     }
 

core/src/main/java/org/elasticsearch/script/NativeScriptFactory.java

@@ -41,4 +41,11 @@ public interface NativeScriptFactory {
      * @param params The parameters passed to the script. Can be <tt>null</tt>.
      */
     ExecutableScript newScript(@Nullable Map<String, Object> params);
+
+    /**
+     * Indicates if document scores may be needed by the produced scripts.
+     * 
+     * @return {@code true} if scores are needed.
+     */
+    boolean needsScores();
 }

core/src/main/java/org/elasticsearch/script/SearchScript.java

@@ -29,4 +29,11 @@ public interface SearchScript {
 
     LeafSearchScript getLeafSearchScript(LeafReaderContext context) throws IOException;
 
+    /**
+     * Indicates if document scores may be needed by this {@link SearchScript}.
+     * 
+     * @return {@code true} if scores are needed.
+     */
+    boolean needsScores();
+
 }

core/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngineService.java

@@ -112,7 +112,6 @@ public class ExpressionScriptEngineService extends AbstractComponent implements
             for (String variable : expr.variables) {
                 if (variable.equals("_score")) {
                     bindings.add(new SortField("_score", SortField.Type.SCORE));
-
                 } else if (variable.equals("_value")) {
                     specialValue = new ReplaceableConstValueSource();
                     bindings.add("_value", specialValue);
@@ -173,7 +172,8 @@ public class ExpressionScriptEngineService extends AbstractComponent implements
                 }
             }
 
-            return new ExpressionSearchScript(compiledScript, bindings, specialValue);
+            final boolean needsScores = expr.getSortField(bindings, false).needsScores();
+            return new ExpressionSearchScript(compiledScript, bindings, specialValue, needsScores);
         } catch (Exception exception) {
             throw new ScriptException("Error during search with " + compiledScript, exception);
         }

core/src/main/java/org/elasticsearch/script/expression/ExpressionSearchScript.java

@@ -46,14 +46,21 @@ class ExpressionSearchScript implements SearchScript {
     final SimpleBindings bindings;
     final ValueSource source;
     final ReplaceableConstValueSource specialValue; // _value
+    final boolean needsScores;
     Scorer scorer;
     int docid;
 
-    ExpressionSearchScript(CompiledScript c, SimpleBindings b, ReplaceableConstValueSource v) {
+    ExpressionSearchScript(CompiledScript c, SimpleBindings b, ReplaceableConstValueSource v, boolean needsScores) {
         compiledScript = c;
         bindings = b;
         source = ((Expression)compiledScript.compiled()).getValueSource(bindings);
         specialValue = v;
+        this.needsScores = needsScores;
+    }
+
+    @Override
+    public boolean needsScores() {
+        return needsScores;
     }
 
     @Override

core/src/main/java/org/elasticsearch/script/groovy/GroovyScriptEngineService.java

@@ -168,6 +168,12 @@ public class GroovyScriptEngineService extends AbstractComponent implements Scri
                 }
                 return new GroovyScript(compiledScript, scriptObject, leafLookup, logger);
             }
+
+            @Override
+            public boolean needsScores() {
+                // TODO: can we reliably know if a groovy script makes use of _score
+                return true;
+            }
         };
     }
 

core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSource.java

@@ -216,8 +216,7 @@ public abstract class ValuesSource {
 
             @Override
             public boolean needsScores() {
-                // TODO: add a way to know whether scripts are using scores
+                return script.needsScores();
-                return true;
             }
         }
 
@@ -295,8 +294,7 @@ public abstract class ValuesSource {
 
             @Override
             public boolean needsScores() {
-                // TODO: add a way to know whether scripts are using scores
+                return script.needsScores();
-                return true;
             }
 
             @Override
@@ -431,8 +429,7 @@ public abstract class ValuesSource {
 
             @Override
             public boolean needsScores() {
-                // TODO: add a way to know whether scripts are using scores
+                return script.needsScores();
-                return true;
             }
         }
 
@@ -451,8 +448,7 @@ public abstract class ValuesSource {
 
         @Override
         public boolean needsScores() {
-            // TODO: add a way to know whether scripts are using scores
+            return script.needsScores();
-            return true;
         }
 
         @Override

core/src/test/java/org/elasticsearch/benchmark/scripts/expression/NativeScript1.java

@@ -34,6 +34,11 @@ public class NativeScript1 extends AbstractSearchScript {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new NativeScript1();
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     public static final String NATIVE_SCRIPT_1 = "native_1";

core/src/test/java/org/elasticsearch/benchmark/scripts/expression/NativeScript2.java

@@ -34,6 +34,11 @@ public class NativeScript2 extends AbstractSearchScript {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new NativeScript2();
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     public static final String NATIVE_SCRIPT_2 = "native_2";

core/src/test/java/org/elasticsearch/benchmark/scripts/expression/NativeScript3.java

@@ -34,6 +34,11 @@ public class NativeScript3 extends AbstractSearchScript {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new NativeScript3();
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     public static final String NATIVE_SCRIPT_3 = "native_3";

core/src/test/java/org/elasticsearch/benchmark/scripts/expression/NativeScript4.java

@@ -34,6 +34,11 @@ public class NativeScript4 extends AbstractSearchScript {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new NativeScript4();
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     public static final String NATIVE_SCRIPT_4 = "native_4";

core/src/test/java/org/elasticsearch/benchmark/scripts/score/script/NativeConstantForLoopScoreScript.java

@@ -36,6 +36,11 @@ public class NativeConstantForLoopScoreScript extends AbstractSearchScript {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new NativeConstantForLoopScoreScript(params);
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     private NativeConstantForLoopScoreScript(Map<String, Object> params) {

core/src/test/java/org/elasticsearch/benchmark/scripts/score/script/NativeConstantScoreScript.java

@@ -36,6 +36,11 @@ public class NativeConstantScoreScript extends AbstractSearchScript {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new NativeConstantScoreScript();
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     private NativeConstantScoreScript() {

core/src/test/java/org/elasticsearch/benchmark/scripts/score/script/NativeNaiveTFIDFScoreScript.java

@@ -42,6 +42,11 @@ public class NativeNaiveTFIDFScoreScript extends AbstractSearchScript {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new NativeNaiveTFIDFScoreScript(params);
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     private NativeNaiveTFIDFScoreScript(Map<String, Object> params) {

core/src/test/java/org/elasticsearch/benchmark/scripts/score/script/NativePayloadSumNoRecordScoreScript.java

@@ -44,6 +44,11 @@ public class NativePayloadSumNoRecordScoreScript extends AbstractSearchScript {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new NativePayloadSumNoRecordScoreScript(params);
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     private NativePayloadSumNoRecordScoreScript(Map<String, Object> params) {

core/src/test/java/org/elasticsearch/benchmark/scripts/score/script/NativePayloadSumScoreScript.java

@@ -44,6 +44,11 @@ public class NativePayloadSumScoreScript extends AbstractSearchScript {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new NativePayloadSumScoreScript(params);
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     private NativePayloadSumScoreScript(Map<String, Object> params) {

core/src/test/java/org/elasticsearch/common/lucene/search/function/ScriptScoreFunctionTests.java

@@ -73,5 +73,10 @@ public class ScriptScoreFunctionTests extends ESTestCase {
                 }
             };
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 }

core/src/test/java/org/elasticsearch/script/NativeScriptTests.java

@@ -97,6 +97,11 @@ public class NativeScriptTests extends ESTestCase {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new MyScript();
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     static class MyScript extends AbstractExecutableScript {

core/src/test/java/org/elasticsearch/script/ScriptFieldIT.java

@@ -81,6 +81,11 @@ public class ScriptFieldIT extends ESIntegTestCase {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new IntScript();
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     static class IntScript extends AbstractSearchScript {
@@ -95,6 +100,11 @@ public class ScriptFieldIT extends ESIntegTestCase {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new LongScript();
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     static class LongScript extends AbstractSearchScript {
@@ -109,6 +119,11 @@ public class ScriptFieldIT extends ESIntegTestCase {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new FloatScript();
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     static class FloatScript extends AbstractSearchScript {
@@ -123,6 +138,11 @@ public class ScriptFieldIT extends ESIntegTestCase {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new DoubleScript();
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     static class DoubleScript extends AbstractSearchScript {

core/src/test/java/org/elasticsearch/script/expression/ExpressionScriptTests.java

@@ -0,0 +1,57 @@
+/*
+ * 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 org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.index.IndexService;
+import org.elasticsearch.script.CompiledScript;
+import org.elasticsearch.script.ScriptService.ScriptType;
+import org.elasticsearch.script.SearchScript;
+import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.test.ESSingleNodeTestCase;
+
+import java.util.Collections;
+
+public class ExpressionScriptTests extends ESSingleNodeTestCase {
+
+    public void testNeedsScores() {
+        IndexService index = createIndex("test", Settings.EMPTY, "type", "d", "type=double");
+        
+        ExpressionScriptEngineService service = new ExpressionScriptEngineService(Settings.EMPTY);
+        SearchLookup lookup = new SearchLookup(index.mapperService(), index.fieldData(), null);
+
+        Object compiled = service.compile("1.2");
+        SearchScript ss = service.search(new CompiledScript(ScriptType.INLINE, "randomName", "expression", compiled), lookup, Collections.<String, Object>emptyMap());
+        assertFalse(ss.needsScores());
+
+        compiled = service.compile("doc['d'].value");
+        ss = service.search(new CompiledScript(ScriptType.INLINE, "randomName", "expression", compiled), lookup, Collections.<String, Object>emptyMap());
+        assertFalse(ss.needsScores());
+
+        compiled = service.compile("1/_score");
+        ss = service.search(new CompiledScript(ScriptType.INLINE, "randomName", "expression", compiled), lookup, Collections.<String, Object>emptyMap());
+        assertTrue(ss.needsScores());
+
+        compiled = service.compile("doc['d'].value * _score");
+        ss = service.search(new CompiledScript(ScriptType.INLINE, "randomName", "expression", compiled), lookup, Collections.<String, Object>emptyMap());
+        assertTrue(ss.needsScores());
+    }
+
+}

core/src/test/java/org/elasticsearch/search/aggregations/bucket/script/NativeSignificanceScoreScriptNoParams.java

@@ -35,6 +35,11 @@ public class NativeSignificanceScoreScriptNoParams extends TestScript {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new NativeSignificanceScoreScriptNoParams();
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     private NativeSignificanceScoreScriptNoParams() {

core/src/test/java/org/elasticsearch/search/aggregations/bucket/script/NativeSignificanceScoreScriptWithParams.java

@@ -36,6 +36,11 @@ public class NativeSignificanceScoreScriptWithParams extends TestScript {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new NativeSignificanceScoreScriptWithParams(params);
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     private NativeSignificanceScoreScriptWithParams(Map<String, Object> params) {

core/src/test/java/org/elasticsearch/search/functionscore/ExplainableScriptIT.java

@@ -102,6 +102,10 @@ public class ExplainableScriptIT extends ESIntegTestCase {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new MyScript();
         }
+        @Override
+        public boolean needsScores() {
+            return true;
+        }
     }
 
     static class MyScript extends AbstractDoubleSearchScript implements ExplainableSearchScript, ExecutableScript {

core/src/test/java/org/elasticsearch/update/UpdateByNativeScriptIT.java

@@ -74,6 +74,10 @@ public class UpdateByNativeScriptIT extends ESIntegTestCase {
         public ExecutableScript newScript(@Nullable Map<String, Object> params) {
             return new CustomScript(params);
         }
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 
     static class CustomScript extends AbstractExecutableScript {

plugins/lang-javascript/src/main/java/org/elasticsearch/script/javascript/JavaScriptScriptEngineService.java

@@ -150,6 +150,12 @@ public class JavaScriptScriptEngineService extends AbstractComponent implements
 
                 return new JavaScriptSearchScript((Script) compiledScript.compiled(), scope, leafLookup);
               }
+
+              @Override
+              public boolean needsScores() {
+                  // TODO: can we reliably know if a javascript script makes use of _score
+                  return true;
+              }
             };
         } finally {
             Context.exit();

plugins/lang-python/src/main/java/org/elasticsearch/script/python/PythonScriptEngineService.java

@@ -90,6 +90,11 @@ public class PythonScriptEngineService extends AbstractComponent implements Scri
                 final LeafSearchLookup leafLookup = lookup.getLeafSearchLookup(context);
                 return new PythonSearchScript((PyCode) compiledScript.compiled(), vars, leafLookup);
             }
+            @Override
+            public boolean needsScores() {
+                // TODO: can we reliably know if a python script makes use of _score
+                return true;
+            }
         };
     }