39
39
import org .apache .lucene .index .DirectoryReader ;
40
40
import org .apache .lucene .index .IndexWriter ;
41
41
import org .apache .lucene .index .LeafReaderContext ;
42
+ import org .apache .lucene .search .DocIdSetIterator ;
42
43
import org .apache .lucene .search .Explanation ;
43
44
import org .apache .lucene .search .IndexSearcher ;
45
+ import org .apache .lucene .search .MatchAllDocsQuery ;
46
+ import org .apache .lucene .search .Query ;
44
47
import org .apache .lucene .search .ScoreMode ;
48
+ import org .apache .lucene .search .Scorer ;
49
+ import org .apache .lucene .search .TwoPhaseIterator ;
45
50
import org .apache .lucene .search .Weight ;
46
51
import org .apache .lucene .store .Directory ;
47
52
import org .opensearch .Version ;
48
53
import org .opensearch .common .lucene .search .Queries ;
49
54
import org .opensearch .common .lucene .search .function .ScriptScoreQuery ;
50
55
import org .opensearch .script .ScoreScript ;
51
56
import org .opensearch .script .Script ;
57
+ import org .opensearch .script .ScriptType ;
52
58
import org .opensearch .search .lookup .LeafSearchLookup ;
53
59
import org .opensearch .search .lookup .SearchLookup ;
54
60
import org .opensearch .test .OpenSearchTestCase ;
55
61
import org .junit .After ;
56
62
import org .junit .Before ;
57
63
58
64
import java .io .IOException ;
65
+ import java .util .HashMap ;
66
+ import java .util .Map ;
59
67
import java .util .function .Function ;
60
68
61
69
import static org .hamcrest .CoreMatchers .containsString ;
@@ -177,6 +185,37 @@ public void testScriptScoreErrorOnNegativeScore() {
177
185
assertTrue (e .getMessage ().contains ("Must be a non-negative score!" ));
178
186
}
179
187
188
+ public void testTwoPhaseIteratorDelegation () throws IOException {
189
+ Map <String , Object > params = new HashMap <>();
190
+ String scriptSource = "doc['field'].value != null ? 2.0 : 0.0" ; // Adjust based on actual field and logic
191
+ Script script = new Script (ScriptType .INLINE , "painless" , scriptSource , params );
192
+ float minScore = 1.0f ; // This should be below the score produced by the script for all docs
193
+ ScoreScript .LeafFactory factory = newFactory (script , false , explanation -> 2.0 );
194
+
195
+ Query subQuery = new MatchAllDocsQuery ();
196
+ ScriptScoreQuery scriptScoreQuery = new ScriptScoreQuery (subQuery , script , factory , minScore , "index" , 0 , Version .CURRENT );
197
+
198
+ Weight weight = searcher .createWeight (searcher .rewrite (scriptScoreQuery ), ScoreMode .COMPLETE , 1f );
199
+
200
+ boolean foundMatchingDoc = false ;
201
+ for (LeafReaderContext leafContext : searcher .getIndexReader ().leaves ()) {
202
+ Scorer scorer = weight .scorer (leafContext );
203
+ if (scorer != null ) {
204
+ TwoPhaseIterator twoPhaseIterator = scorer .twoPhaseIterator ();
205
+ assertNotNull ("TwoPhaseIterator should not be null" , twoPhaseIterator );
206
+ DocIdSetIterator docIdSetIterator = twoPhaseIterator .approximation ();
207
+ int docId ;
208
+ while ((docId = docIdSetIterator .nextDoc ()) != DocIdSetIterator .NO_MORE_DOCS ) {
209
+ if (twoPhaseIterator .matches ()) {
210
+ foundMatchingDoc = true ;
211
+ break ;
212
+ }
213
+ }
214
+ }
215
+ }
216
+ assertTrue ("Expected to find at least one matching document" , foundMatchingDoc );
217
+ }
218
+
180
219
private ScoreScript .LeafFactory newFactory (
181
220
Script script ,
182
221
boolean needsScore ,
@@ -203,5 +242,4 @@ public double execute(ExplanationHolder explanation) {
203
242
}
204
243
};
205
244
}
206
-
207
245
}
0 commit comments