Skip to content

Commit 684dd74

Browse files
[Bugfix] Pass TwoPhaseIterator from subQueryScorer (opensearch-project#11954) (opensearch-project#12213)
(cherry picked from commit 0d50525) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
1 parent da3c1a5 commit 684dd74

File tree

2 files changed

+45
-1
lines changed

2 files changed

+45
-1
lines changed

server/src/main/java/org/opensearch/common/lucene/search/function/ScriptScoreQuery.java

+6
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import org.apache.lucene.search.Scorable;
4646
import org.apache.lucene.search.ScoreMode;
4747
import org.apache.lucene.search.Scorer;
48+
import org.apache.lucene.search.TwoPhaseIterator;
4849
import org.apache.lucene.search.Weight;
4950
import org.apache.lucene.util.Bits;
5051
import org.opensearch.Version;
@@ -302,6 +303,11 @@ public DocIdSetIterator iterator() {
302303
return subQueryScorer.iterator();
303304
}
304305

306+
@Override
307+
public TwoPhaseIterator twoPhaseIterator() {
308+
return subQueryScorer.twoPhaseIterator();
309+
}
310+
305311
@Override
306312
public float getMaxScore(int upTo) {
307313
return Float.MAX_VALUE; // TODO: what would be a good upper bound?

server/src/test/java/org/opensearch/search/query/ScriptScoreQueryTests.java

+39-1
Original file line numberDiff line numberDiff line change
@@ -39,23 +39,31 @@
3939
import org.apache.lucene.index.DirectoryReader;
4040
import org.apache.lucene.index.IndexWriter;
4141
import org.apache.lucene.index.LeafReaderContext;
42+
import org.apache.lucene.search.DocIdSetIterator;
4243
import org.apache.lucene.search.Explanation;
4344
import org.apache.lucene.search.IndexSearcher;
45+
import org.apache.lucene.search.MatchAllDocsQuery;
46+
import org.apache.lucene.search.Query;
4447
import org.apache.lucene.search.ScoreMode;
48+
import org.apache.lucene.search.Scorer;
49+
import org.apache.lucene.search.TwoPhaseIterator;
4550
import org.apache.lucene.search.Weight;
4651
import org.apache.lucene.store.Directory;
4752
import org.opensearch.Version;
4853
import org.opensearch.common.lucene.search.Queries;
4954
import org.opensearch.common.lucene.search.function.ScriptScoreQuery;
5055
import org.opensearch.script.ScoreScript;
5156
import org.opensearch.script.Script;
57+
import org.opensearch.script.ScriptType;
5258
import org.opensearch.search.lookup.LeafSearchLookup;
5359
import org.opensearch.search.lookup.SearchLookup;
5460
import org.opensearch.test.OpenSearchTestCase;
5561
import org.junit.After;
5662
import org.junit.Before;
5763

5864
import java.io.IOException;
65+
import java.util.HashMap;
66+
import java.util.Map;
5967
import java.util.function.Function;
6068

6169
import static org.hamcrest.CoreMatchers.containsString;
@@ -177,6 +185,37 @@ public void testScriptScoreErrorOnNegativeScore() {
177185
assertTrue(e.getMessage().contains("Must be a non-negative score!"));
178186
}
179187

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+
180219
private ScoreScript.LeafFactory newFactory(
181220
Script script,
182221
boolean needsScore,
@@ -203,5 +242,4 @@ public double execute(ExplanationHolder explanation) {
203242
}
204243
};
205244
}
206-
207245
}

0 commit comments

Comments
 (0)