Fix LTR query feature with phrases (and two-phase) queries

Query features should verify that docs match the two-phase iterator.

Author: jimczi

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/ltr/QueryFeatureExtractor.java

@@ -15,7 +15,6 @@
 import org.apache.lucene.search.Weight;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
@@ -25,52 +24,52 @@
  * respective feature name.
  */
 public class QueryFeatureExtractor implements FeatureExtractor {
-
     private final List<String> featureNames;
     private final List<Weight> weights;
-    private final List<Scorer> scorers;
-    private DisjunctionDISIApproximation rankerIterator;
+
+    private final DisiPriorityQueue subScorers;
+    private DisjunctionDISIApproximation approximation;
 
     public QueryFeatureExtractor(List<String> featureNames, List<Weight> weights) {
         if (featureNames.size() != weights.size()) {
             throw new IllegalArgumentException("[featureNames] and [weights] must be the same size.");
         }
         this.featureNames = featureNames;
         this.weights = weights;
-        this.scorers = new ArrayList<>(weights.size());
+        this.subScorers = new DisiPriorityQueue(weights.size());
     }
 
     @Override
     public void setNextReader(LeafReaderContext segmentContext) throws IOException {
-        DisiPriorityQueue disiPriorityQueue = new DisiPriorityQueue(weights.size());
-        scorers.clear();
-        for (Weight weight : weights) {
+        subScorers.clear();
+        for (int i = 0; i < weights.size(); i++) {
+            var weight = weights.get(i);
             if (weight == null) {
-                scorers.add(null);
                 continue;
             }
             Scorer scorer = weight.scorer(segmentContext);
             if (scorer != null) {
-                disiPriorityQueue.add(new DisiWrapper(scorer, false));
+                subScorers.add(new FeatureDisiWrapper(scorer, featureNames.get(i)));
             }
-            scorers.add(scorer);
         }
-
-        rankerIterator = disiPriorityQueue.size() > 0 ? new DisjunctionDISIApproximation(disiPriorityQueue) : null;
+        approximation = subScorers.size() > 0 ? new DisjunctionDISIApproximation(subScorers) : null;
     }
 
     @Override
     public void addFeatures(Map<String, Object> featureMap, int docId) throws IOException {
-        if (rankerIterator == null) {
+        if (approximation == null || approximation.docID() > docId) {
             return;
         }
-
-        rankerIterator.advance(docId);
-        for (int i = 0; i < featureNames.size(); i++) {
-            Scorer scorer = scorers.get(i);
-            // Do we have a scorer, and does it match the provided document?
-            if (scorer != null && scorer.docID() == docId) {
-                featureMap.put(featureNames.get(i), scorer.score());
+        if (approximation.docID() < docId) {
+            approximation.advance(docId);
+        }
+        if (approximation.docID() != docId) {
+            return;
+        }
+        var w = (FeatureDisiWrapper) subScorers.topList();
+        for (; w != null; w = (FeatureDisiWrapper) w.next) {
+            if (w.twoPhaseView == null || w.twoPhaseView.matches()) {
+                featureMap.put(w.featureName, w.scorable.score());
             }
         }
     }
@@ -80,4 +79,11 @@ public List<String> featureNames() {
         return featureNames;
     }
 
+    private static class FeatureDisiWrapper extends DisiWrapper {
+        final String featureName;
+        FeatureDisiWrapper(Scorer scorer, String featureName) {
+            super(scorer, false);
+            this.featureName = featureName;
+        }
+    }
 }

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/inference/ltr/QueryFeatureExtractorTests.java

@@ -65,7 +65,6 @@ private void addDocs(String[] textValues, int[] numberValues) throws IOException
         searcher.setSimilarity(new ClassicSimilarity());
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/98127")
     public void testQueryExtractor() throws IOException {
         addDocs(
             new String[] { "the quick brown fox", "the slow brown fox", "the grey dog", "yet another string" },
@@ -86,6 +85,8 @@ public void testQueryExtractor() throws IOException {
             new QueryExtractorBuilder(
                 "matching_missing_field",
                 QueryProvider.fromParsedQuery(QueryBuilders.termQuery("missing_text", "quick fox"))
+            ).rewrite(ctx),
+            new QueryExtractorBuilder("phrase_score", QueryProvider.fromParsedQuery(QueryBuilders.matchPhraseQuery(TEXT_FIELD_NAME, "slow brown fox"))
             ).rewrite(ctx)
         );
         SearchExecutionContext dummySEC = createSearchExecutionContext();
@@ -117,11 +118,15 @@ public void testQueryExtractor() throws IOException {
         // First two only match the text field
         assertThat(extractedFeatures.get(0), hasEntry("text_score", 1.7135582f));
         assertThat(extractedFeatures.get(0), not(hasKey("number_score")));
+        assertThat(extractedFeatures.get(0), not(hasKey("phrase_score")));
         assertThat(extractedFeatures.get(1), hasEntry("text_score", 0.7554128f));
         assertThat(extractedFeatures.get(1), not(hasKey("number_score")));
+        assertThat(extractedFeatures.get(1), hasEntry("phrase_score", 2.468971f));
+
         // Only matches the range query
         assertThat(extractedFeatures.get(2), hasEntry("number_score", 1f));
         assertThat(extractedFeatures.get(2), not(hasKey("text_score")));
+        assertThat(extractedFeatures.get(2), not(hasKey("phrase_score")));
         // No query matches
         assertThat(extractedFeatures.get(3), anEmptyMap());
         reader.close();