FunctionScore: Fix 'avg' score mode to correctly implement weighted m…

…ean. closes #8992 closes #9004
elastic · Dec 19, 2014 · 8738583 · 8738583
1 parent e6a190e
commit 8738583
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 9 deletions.
diff --git a/src/main/java/org/elasticsearch/common/lucene/search/function/FiltersFunctionScoreQuery.java b/src/main/java/org/elasticsearch/common/lucene/search/function/FiltersFunctionScoreQuery.java
@@ -177,7 +177,15 @@ public Explanation explain(LeafReaderContext context, int doc) throws IOExceptio
             }
             // First: Gather explanations for all filters
             List<ComplexExplanation> filterExplanations = new ArrayList<>();
+            float weightSum = 0;
             for (FilterFunction filterFunction : filterFunctions) {
+
+                if (filterFunction.function instanceof WeightFactorFunction) {
+                    weightSum += ((WeightFactorFunction) filterFunction.function).getWeight();
+                } else {
+                    weightSum++;
+                }
+
                 Bits docSet = DocIdSets.toSafeBits(context.reader(),
                         filterFunction.filter.getDocIdSet(context, context.reader().getLiveDocs()));
                 if (docSet.get(doc)) {
@@ -226,15 +234,13 @@ public Explanation explain(LeafReaderContext context, int doc) throws IOExceptio
                 break;
             default: // Avg / Total
                 double totalFactor = 0.0f;
-                int count = 0;
                 for (int i = 0; i < filterExplanations.size(); i++) {
                     totalFactor += filterExplanations.get(i).getValue();
-                    count++;
                 }
-                if (count != 0) {
+                if (weightSum != 0) {
                     factor = totalFactor;
                     if (scoreMode == ScoreMode.Avg) {
-                        factor /= count;
+                        factor /= weightSum;
                     }
                 }
             }
@@ -300,17 +306,21 @@ public float innerScore() throws IOException {
                 }
             } else { // Avg / Total
                 double totalFactor = 0.0f;
-                int count = 0;
+                float weightSum = 0;
                 for (int i = 0; i < filterFunctions.length; i++) {
                     if (docSets[i].get(docId)) {
                         totalFactor += filterFunctions[i].function.score(docId, subQueryScore);
-                        count++;
+                        if (filterFunctions[i].function instanceof WeightFactorFunction) {
+                            weightSum+= ((WeightFactorFunction)filterFunctions[i].function).getWeight();
+                        } else {
+                            weightSum++;
+                        }
                     }
                 }
-                if (count != 0) {
+                if (weightSum != 0) {
                     factor = totalFactor;
                     if (scoreMode == ScoreMode.Avg) {
-                        factor /= count;
+                        factor /= weightSum;
                     }
                 }
             }

diff --git a/src/test/java/org/elasticsearch/search/functionscore/FunctionScoreTests.java b/src/test/java/org/elasticsearch/search/functionscore/FunctionScoreTests.java
@@ -254,8 +254,11 @@ protected double computeExpectedScore(float[] weights, float[] scores, String sc
             expectedScore = Float.MAX_VALUE;
         }
 
+        float weightSum = 0;
+
         for (int i = 0; i < weights.length; i++) {
             double functionScore = (double) weights[i] * scores[i];
+            weightSum += weights[i];
 
             if ("avg".equals(scoreMode)) {
                 expectedScore += functionScore;
@@ -271,7 +274,7 @@ protected double computeExpectedScore(float[] weights, float[] scores, String sc
 
         }
         if ("avg".equals(scoreMode)) {
-            expectedScore /= weights.length;
+            expectedScore /= weightSum;
         }
         return expectedScore;
     }