Merge pull request #12693 from jpountz/enhancement/speed_up_function_score_query_when_replacing_score

Speed up the `function_score` query when scores are not needed.

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

@@ -57,6 +57,11 @@ public class BoostScoreFunction extends ScoreFunction {
         };
     }
 
+    @Override
+    public boolean needsScores() {
+        return false;
+    }
+
     @Override
     public String toString() {
         return "boost[" + boost + "]";

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

@@ -91,6 +91,11 @@ public class FieldValueFactorFunction extends ScoreFunction {
         };
     }
 
+    @Override
+    public boolean needsScores() {
+        return false;
+    }
+
     /**
      * The Type class encapsulates the modification types that can be applied
      * to the score/value product.

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

@@ -44,7 +44,7 @@ public class FunctionScoreQuery extends Query {
     public FunctionScoreQuery(Query subQuery, ScoreFunction function, Float minScore) {
         this.subQuery = subQuery;
         this.function = function;
-        this.combineFunction = function == null? combineFunction.MULT : function.getDefaultScoreCombiner();
+        this.combineFunction = function == null? CombineFunction.MULT : function.getDefaultScoreCombiner();
         this.minScore = minScore;
     }
 
@@ -87,19 +87,27 @@ public class FunctionScoreQuery extends Query {
 
     @Override
     public Weight createWeight(IndexSearcher searcher, boolean needsScores) throws IOException {
-        // TODO: needsScores
-        // if we don't need scores, just return the underlying weight?
-        Weight subQueryWeight = subQuery.createWeight(searcher, needsScores);
-        return new CustomBoostFactorWeight(this, subQueryWeight);
+        if (needsScores == false) {
+            return subQuery.createWeight(searcher, needsScores);
+        }
+
+        boolean subQueryNeedsScores =
+                combineFunction != CombineFunction.REPLACE // if we don't replace we need the original score
+                || function == null // when the function is null, we just multiply the score, so we need it
+                || function.needsScores(); // some scripts can replace with a script that returns eg. 1/_score
+        Weight subQueryWeight = subQuery.createWeight(searcher, subQueryNeedsScores);
+        return new CustomBoostFactorWeight(this, subQueryWeight, subQueryNeedsScores);
     }
 
     class CustomBoostFactorWeight extends Weight {
 
         final Weight subQueryWeight;
+        final boolean needsScores;
 
-        public CustomBoostFactorWeight(Query parent, Weight subQueryWeight) throws IOException {
+        public CustomBoostFactorWeight(Query parent, Weight subQueryWeight, boolean needsScores) throws IOException {
             super(parent);
             this.subQueryWeight = subQueryWeight;
+            this.needsScores = needsScores;
         }
 
         @Override
@@ -129,7 +137,7 @@ public class FunctionScoreQuery extends Query {
             if (function != null) {
                 leafFunction = function.getLeafScoreFunction(context);
             }
-            return new FunctionFactorScorer(this, subQueryScorer, leafFunction, maxBoost, combineFunction, minScore);
+            return new FunctionFactorScorer(this, subQueryScorer, leafFunction, maxBoost, combineFunction, minScore, needsScores);
         }
 
         @Override
@@ -150,16 +158,21 @@ public class FunctionScoreQuery extends Query {
     static class FunctionFactorScorer extends CustomBoostFactorScorer {
 
         private final LeafScoreFunction function;
+        private final boolean needsScores;
 
-        private FunctionFactorScorer(CustomBoostFactorWeight w, Scorer scorer, LeafScoreFunction function, float maxBoost, CombineFunction scoreCombiner, Float minScore)
+        private FunctionFactorScorer(CustomBoostFactorWeight w, Scorer scorer, LeafScoreFunction function, float maxBoost, CombineFunction scoreCombiner, Float minScore, boolean needsScores)
                 throws IOException {
             super(w, scorer, maxBoost, scoreCombiner, minScore);
             this.function = function;
+            this.needsScores = needsScores;
         }
 
         @Override
         public float innerScore() throws IOException {
-            float score = scorer.score();
+            // Even if the weight is created with needsScores=false, it might
+            // be costly to call score(), so we explicitly check if scores
+            // are needed
+            float score = needsScores ? scorer.score() : 0f;
             if (function == null) {
                 return subQueryBoost * score;
             } else {

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

@@ -81,4 +81,8 @@ public class RandomScoreFunction extends ScoreFunction {
         };
     }
 
+    @Override
+    public boolean needsScores() {
+        return false;
+    }
 }

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

@@ -39,4 +39,11 @@ public abstract class ScoreFunction {
     }
 
     public abstract LeafScoreFunction getLeafScoreFunction(LeafReaderContext ctx) throws IOException;
+
+    /**
+     * Indicates if document scores are needed by this function.
+     * 
+     * @return {@code true} if scores are needed.
+     */
+    public abstract boolean needsScores();
 }

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

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

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

@@ -71,6 +71,11 @@ public class WeightFactorFunction extends ScoreFunction {
         };
     }
 
+    @Override
+    public boolean needsScores() {
+        return false;
+    }
+
     public Explanation explainWeight() {
         return Explanation.match(getWeight(), "weight");
     }
@@ -99,5 +104,10 @@ public class WeightFactorFunction extends ScoreFunction {
                 }
             };
         }
+
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
     }
 }

core/src/main/java/org/elasticsearch/index/query/functionscore/DecayFunctionParser.java

@@ -289,6 +289,11 @@ public abstract class DecayFunctionParser implements ScoreFunctionParser {
             this.fieldData = fieldData;
         }
 
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
+
         @Override
         protected NumericDoubleValues distance(LeafReaderContext context) {
             final MultiGeoPointValues geoPointValues = fieldData.load(context).getGeoPointValues();
@@ -352,6 +357,11 @@ public abstract class DecayFunctionParser implements ScoreFunctionParser {
             this.origin = origin;
         }
 
+        @Override
+        public boolean needsScores() {
+            return false;
+        }
+
         @Override
         protected NumericDoubleValues distance(LeafReaderContext context) {
             final SortedNumericDoubleValues doubleValues = fieldData.load(context).getDoubleValues();