percolator: Also support query term extract for queries wrapped inside a FunctionScoreQuery

Additionally for highlighting percolator hits, also extract percolator query from FunctionScoreQuery and DisjunctionMaxQuery

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

@@ -73,6 +73,10 @@ public class FunctionScoreQuery extends Query {
         return function;
     }
 
+    public Float getMinScore() {
+        return minScore;
+    }
+
     @Override
     public Query rewrite(IndexReader reader) throws IOException {
         Query rewritten = super.rewrite(reader);

modules/percolator/src/main/java/org/elasticsearch/percolator/ExtractQueryTermsService.java

@@ -49,6 +49,7 @@ import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.BytesRefBuilder;
 import org.elasticsearch.common.logging.LoggerMessageFormat;
 import org.elasticsearch.common.lucene.search.MatchNoDocsQuery;
+import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
 import org.elasticsearch.index.mapper.ParseContext;
 
 import java.io.IOException;
@@ -75,7 +76,7 @@ public final class ExtractQueryTermsService {
     static final Map<Class<? extends Query>, Function<Query, Result>> queryProcessors;
 
     static {
-        Map<Class<? extends Query>, Function<Query, Result>> map = new HashMap<>(16);
+        Map<Class<? extends Query>, Function<Query, Result>> map = new HashMap<>(17);
         map.put(MatchNoDocsQuery.class, matchNoDocsQuery());
         map.put(ConstantScoreQuery.class, constantScoreQuery());
         map.put(BoostQuery.class, boostQuery());
@@ -92,6 +93,7 @@ public final class ExtractQueryTermsService {
         map.put(BooleanQuery.class, booleanQuery());
         map.put(DisjunctionMaxQuery.class, disjunctionMaxQuery());
         map.put(SynonymQuery.class, synonymQuery());
+        map.put(FunctionScoreQuery.class, functionScoreQuery());
         queryProcessors = Collections.unmodifiableMap(map);
     }
 
@@ -376,6 +378,19 @@ public final class ExtractQueryTermsService {
         };
     }
 
+    static Function<Query, Result> functionScoreQuery() {
+        return query -> {
+            FunctionScoreQuery functionScoreQuery = (FunctionScoreQuery) query;
+            Result result = extractQueryTerms(functionScoreQuery.getSubQuery());
+            // If min_score is specified we can't guarantee upfront that this percolator query matches,
+            // so in that case we set verified to false.
+            // (if it matches with the percolator document matches with the extracted terms.
+            // Min score filters out docs, which is different than the functions, which just influences the score.)
+            boolean verified = functionScoreQuery.getMinScore() == null;
+            return new Result(verified, result.terms);
+        };
+    }
+
     static Result handleDisjunction(List<Query> disjunctions, int minimumShouldMatch, boolean otherClauses) {
         boolean verified = minimumShouldMatch <= 1 && otherClauses == false;
         Set<Term> terms = new HashSet<>();

modules/percolator/src/main/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhase.java

@@ -25,9 +25,11 @@ import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.BoostQuery;
 import org.apache.lucene.search.ConstantScoreQuery;
+import org.apache.lucene.search.DisjunctionMaxQuery;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.text.Text;
 import org.elasticsearch.index.query.ParsedQuery;
@@ -110,10 +112,19 @@ public final class PercolatorHighlightSubFetchPhase extends HighlightPhase {
                     return result;
                 }
             }
+        } else if (query instanceof DisjunctionMaxQuery) {
+            for (Query disjunct : ((DisjunctionMaxQuery) query).getDisjuncts()) {
+                PercolateQuery result = locatePercolatorQuery(disjunct);
+                if (result != null) {
+                    return result;
+                }
+            }
         } else if (query instanceof ConstantScoreQuery) {
             return locatePercolatorQuery(((ConstantScoreQuery) query).getQuery());
         } else if (query instanceof BoostQuery) {
             return locatePercolatorQuery(((BoostQuery) query).getQuery());
+        } else if (query instanceof FunctionScoreQuery) {
+            return locatePercolatorQuery(((FunctionScoreQuery) query).getSubQuery());
         }
 
         return null;

modules/percolator/src/test/java/org/elasticsearch/percolator/ExtractQueryTermsServiceTests.java

@@ -46,6 +46,8 @@ import org.apache.lucene.search.spans.SpanOrQuery;
 import org.apache.lucene.search.spans.SpanTermQuery;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.lucene.search.MatchNoDocsQuery;
+import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
+import org.elasticsearch.common.lucene.search.function.RandomScoreFunction;
 import org.elasticsearch.index.mapper.ParseContext;
 import org.elasticsearch.percolator.ExtractQueryTermsService.Result;
 import org.elasticsearch.test.ESTestCase;
@@ -571,6 +573,19 @@ public class ExtractQueryTermsServiceTests extends ESTestCase {
         assertTermsEqual(result.terms, new Term("_field", "_value1"), new Term("_field", "_value2"));
     }
 
+    public void testFunctionScoreQuery() {
+        TermQuery termQuery = new TermQuery(new Term("_field", "_value"));
+        FunctionScoreQuery functionScoreQuery = new FunctionScoreQuery(termQuery, new RandomScoreFunction());
+        Result result = extractQueryTerms(functionScoreQuery);
+        assertThat(result.verified, is(true));
+        assertTermsEqual(result.terms, new Term("_field", "_value"));
+
+        functionScoreQuery = new FunctionScoreQuery(termQuery, new RandomScoreFunction(), 1f, null, 10f);
+        result = extractQueryTerms(functionScoreQuery);
+        assertThat(result.verified, is(false));
+        assertTermsEqual(result.terms, new Term("_field", "_value"));
+    }
+
     public void testCreateQueryMetadataQuery() throws Exception {
         MemoryIndex memoryIndex = new MemoryIndex(false);
         memoryIndex.addField("field1", "the quick brown fox jumps over the lazy dog", new WhitespaceAnalyzer());

modules/percolator/src/test/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhaseTests.java

@@ -22,9 +22,12 @@ import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.BoostQuery;
 import org.apache.lucene.search.ConstantScoreQuery;
+import org.apache.lucene.search.DisjunctionMaxQuery;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.elasticsearch.common.bytes.BytesArray;
+import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
+import org.elasticsearch.common.lucene.search.function.RandomScoreFunction;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.search.Highlighters;
 import org.elasticsearch.search.highlight.SearchContextHighlight;
@@ -32,6 +35,7 @@ import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.test.ESTestCase;
 import org.mockito.Mockito;
 
+import java.util.Arrays;
 import java.util.Collections;
 
 import static org.hamcrest.Matchers.is;
@@ -77,6 +81,16 @@ public class PercolatorHighlightSubFetchPhaseTests extends ESTestCase {
         assertThat(PercolatorHighlightSubFetchPhase.locatePercolatorQuery(boostQuery), nullValue());
         boostQuery = new BoostQuery(percolateQuery, 1f);
         assertThat(PercolatorHighlightSubFetchPhase.locatePercolatorQuery(boostQuery), sameInstance(percolateQuery));
+
+        FunctionScoreQuery functionScoreQuery = new FunctionScoreQuery(new MatchAllDocsQuery(), new RandomScoreFunction());
+        assertThat(PercolatorHighlightSubFetchPhase.locatePercolatorQuery(functionScoreQuery), nullValue());
+        functionScoreQuery = new FunctionScoreQuery(percolateQuery, new RandomScoreFunction());
+        assertThat(PercolatorHighlightSubFetchPhase.locatePercolatorQuery(functionScoreQuery), sameInstance(percolateQuery));
+
+        DisjunctionMaxQuery disjunctionMaxQuery = new DisjunctionMaxQuery(Arrays.asList(new MatchAllDocsQuery()), 1f);
+        assertThat(PercolatorHighlightSubFetchPhase.locatePercolatorQuery(disjunctionMaxQuery), nullValue());
+        disjunctionMaxQuery = new DisjunctionMaxQuery(Arrays.asList(percolateQuery, new MatchAllDocsQuery()), 1f);
+        assertThat(PercolatorHighlightSubFetchPhase.locatePercolatorQuery(disjunctionMaxQuery), sameInstance(percolateQuery));
     }
 
 }