Leverage scorer supplier in QueryFeatureExtractor

jimczi · jimczi · commit c3a8271f6db2 · 2025-03-19T20:35:59.000Z
Follow up of elastic#125103 that leverages scorer supplier to create queries optimised to run on top docs only.
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/ltr/QueryFeatureExtractor.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/ltr/QueryFeatureExtractor.java
@@ -47,9 +47,12 @@ public void setNextReader(LeafReaderContext segmentContext) throws IOException {
             if (weight == null) {
                 continue;
             }
-            Scorer scorer = weight.scorer(segmentContext);
-            if (scorer != null) {
-                subScorers.add(new FeatureDisiWrapper(scorer, featureNames.get(i)));
+            var scorerSupplier = weight.scorerSupplier(segmentContext);
+            if (scorerSupplier != null) {
+                var scorer = scorerSupplier.get(0L);
+                if (scorer != null) {
+                    subScorers.add(new FeatureDisiWrapper(scorer, featureNames.get(i)));
+                }
             }
         }
         approximation = subScorers.size() > 0 ? new DisjunctionDISIApproximation(subScorers) : null;

Original file line number	Diff line number	Diff line change
`@@ -47,9 +47,12 @@ public void setNextReader(LeafReaderContext segmentContext) throws IOException {`
`47`	`47`	`if (weight == null) {`
`48`	`48`	`continue;`
`49`	`49`	`}`
`50`		`- Scorer scorer = weight.scorer(segmentContext);`
`51`		`- if (scorer != null) {`
`52`		`- subScorers.add(new FeatureDisiWrapper(scorer, featureNames.get(i)));`
	`50`	`+ var scorerSupplier = weight.scorerSupplier(segmentContext);`
	`51`	`+ if (scorerSupplier != null) {`
	`52`	`+ var scorer = scorerSupplier.get(0L);`
	`53`	`+ if (scorer != null) {`
	`54`	`+ subScorers.add(new FeatureDisiWrapper(scorer, featureNames.get(i)));`
	`55`	`+ }`
`53`	`56`	`}`
`54`	`57`	`}`
`55`	`58`	`approximation = subScorers.size() > 0 ? new DisjunctionDISIApproximation(subScorers) : null;`