@Test
public void testMultipleRescores() throws Exception {
int numDocs = indexRandomNumbers("keyword", 1);
QueryRescorer eightIsGreat =
RescoreBuilder.queryRescorer(
QueryBuilders.functionScoreQuery(
QueryBuilders.termQuery("field1", English.intToEnglish(8)))
.boostMode(CombineFunction.REPLACE)
.add(ScoreFunctionBuilders.scriptFunction("1000.0f")))
.setScoreMode("total");
QueryRescorer sevenIsBetter =
RescoreBuilder.queryRescorer(
QueryBuilders.functionScoreQuery(
QueryBuilders.termQuery("field1", English.intToEnglish(7)))
.boostMode(CombineFunction.REPLACE)
.add(ScoreFunctionBuilders.scriptFunction("10000.0f")))
.setScoreMode("total");
// First set the rescore window large enough that both rescores take effect
SearchRequestBuilder request = client().prepareSearch().setRescoreWindow(numDocs);
request.addRescorer(eightIsGreat).addRescorer(sevenIsBetter);
SearchResponse response = request.get();
assertFirstHit(response, hasId("7"));
assertSecondHit(response, hasId("8"));
// Now squash the second rescore window so it never gets to see a seven
response =
request
.setSize(1)
.clearRescorers()
.addRescorer(eightIsGreat)
.addRescorer(sevenIsBetter, 1)
.get();
assertFirstHit(response, hasId("8"));
// We have no idea what the second hit will be because we didn't get a chance to look for seven
// Now use one rescore to drag the number we're looking for into the window of another
QueryRescorer ninetyIsGood =
RescoreBuilder.queryRescorer(
QueryBuilders.functionScoreQuery(QueryBuilders.queryString("*ninety*"))
.boostMode(CombineFunction.REPLACE)
.add(ScoreFunctionBuilders.scriptFunction("1000.0f")))
.setScoreMode("total");
QueryRescorer oneToo =
RescoreBuilder.queryRescorer(
QueryBuilders.functionScoreQuery(QueryBuilders.queryString("*one*"))
.boostMode(CombineFunction.REPLACE)
.add(ScoreFunctionBuilders.scriptFunction("1000.0f")))
.setScoreMode("total");
request.clearRescorers().addRescorer(ninetyIsGood).addRescorer(oneToo, 10);
response = request.setSize(2).get();
assertFirstHit(response, hasId("91"));
assertFirstHit(response, hasScore(2001.0f));
assertSecondHit(
response, hasScore(1001.0f)); // Not sure which one it is but it is ninety something
}
private void checkValueInEachDocWithFunctionScore(
String fieldScript,
Map<String, Object> expectedFieldVals,
String scoreScript,
Map<String, Object> expectedScore,
int numExpectedDocs) {
SearchResponse sr =
client()
.prepareSearch("test")
.setQuery(
QueryBuilders.functionScoreQuery(ScoreFunctionBuilders.scriptFunction(scoreScript)))
.addScriptField("tvtest", fieldScript)
.execute()
.actionGet();
assertHitCount(sr, numExpectedDocs);
for (SearchHit hit : sr.getHits().getHits()) {
Object result = hit.getFields().get("tvtest").getValues().get(0);
Object expectedResult = expectedFieldVals.get(hit.getId());
assertThat("for doc " + hit.getId(), result, equalTo(expectedResult));
assertThat(
"for doc " + hit.getId(),
((Float) expectedScore.get(hit.getId())).doubleValue(),
Matchers.closeTo(hit.score(), 1.e-4));
}
}
// Query based on term
public static HashMap<String, Double> queryTF(
Client client,
String qb,
String index,
String type,
HashMap<String, Double> result,
HashMap<String, Double> lengthmap,
HashMap<String, Double> staticmap)
throws IOException, JSONException {
HashMap parammap = new HashMap();
parammap.put("field", "text");
parammap.put("term", qb);
SearchResponse responsesearch =
client
.prepareSearch(index)
.setQuery(
new FunctionScoreQueryBuilder(QueryBuilders.matchQuery("text", qb))
.boostMode("replace")
.add(ScoreFunctionBuilders.scriptFunction("getTF").params(parammap)))
.setSize(100000)
.setNoFields()
.execute()
.actionGet();
JSONObject obj = new JSONObject(responsesearch);
JSONObject obj2 = obj.getJSONObject("hits");
JSONArray hits = obj2.getJSONArray("hits");
Iterator itrnext = staticmap.keySet().iterator();
while (itrnext.hasNext()) {
String id = itrnext.next().toString();
double oldtf = staticmap.get(id);
double v = 178081;
double len = lengthmap.get(id);
double newtf = 1 / (len + v);
double lap = Math.log(newtf);
staticmap.put(id, oldtf + lap);
}
for (int i = 0; i < hits.length(); i++) {
JSONObject newobj = hits.getJSONObject(i);
String id = newobj.getString("id");
Double tf = newobj.getDouble("score");
double len = lengthmap.get(id);
double v = 178081;
double lap = (tf + 1) / (len + v);
double inc = 1 / (len + v);
lap = Math.log(lap);
inc = Math.log(inc);
if (staticmap.containsKey(id)) {
double oldtf = staticmap.get(id);
staticmap.put(id, oldtf + lap - inc);
} else staticmap.put(id, lap);
}
return staticmap;
}
private void checkOnlyFunctionScore(
String scoreScript, Map<String, Object> expectedScore, int numExpectedDocs) {
SearchResponse sr =
client()
.prepareSearch("test")
.setQuery(
QueryBuilders.functionScoreQuery(ScoreFunctionBuilders.scriptFunction(scoreScript)))
.execute()
.actionGet();
assertHitCount(sr, numExpectedDocs);
for (SearchHit hit : sr.getHits().getHits()) {
assertThat(
"for doc " + hit.getId(),
((Float) expectedScore.get(hit.getId())).doubleValue(),
Matchers.closeTo(hit.score(), 1.e-4));
}
}
@Test
public void testEnforceWindowSize() {
createIndex("test");
// this
int iters = atLeast(10);
for (int i = 0; i < iters; i++) {
client()
.prepareIndex("test", "type", Integer.toString(i))
.setSource("f", Integer.toString(i))
.execute()
.actionGet();
}
refresh();
int numShards = getNumShards("test").numPrimaries;
for (int j = 0; j < iters; j++) {
SearchResponse searchResponse =
client()
.prepareSearch()
.setQuery(QueryBuilders.matchAllQuery())
.setRescorer(
RescoreBuilder.queryRescorer(
QueryBuilders.functionScoreQuery(QueryBuilders.matchAllQuery())
.boostMode("replace")
.add(ScoreFunctionBuilders.factorFunction(100)))
.setQueryWeight(0.0f)
.setRescoreQueryWeight(1.0f))
.setRescoreWindow(1)
.setSize(randomIntBetween(2, 10))
.execute()
.actionGet();
assertFirstHit(searchResponse, hasScore(100.f));
int numPending100 = numShards;
for (int i = 0; i < searchResponse.getHits().hits().length; i++) {
float score = searchResponse.getHits().hits()[i].getScore();
if (score == 100f) {
assertThat(numPending100--, greaterThanOrEqualTo(0));
} else {
assertThat(numPending100, equalTo(0));
}
}
}
}
public void testScore() throws Exception {
createIndex("test");
ensureGreen("test");
indexRandom(
true,
client().prepareIndex("test", "doc", "1").setSource("text", "hello goodbye"),
client().prepareIndex("test", "doc", "2").setSource("text", "hello hello hello goodbye"),
client().prepareIndex("test", "doc", "3").setSource("text", "hello hello goodebye"));
ScoreFunctionBuilder score =
ScoreFunctionBuilders.scriptFunction(
new Script("1 / _score", ScriptType.INLINE, "expression", null));
SearchRequestBuilder req = client().prepareSearch().setIndices("test");
req.setQuery(
QueryBuilders.functionScoreQuery(QueryBuilders.termQuery("text", "hello"), score)
.boostMode("replace"));
req.setSearchType(SearchType.DFS_QUERY_THEN_FETCH); // make sure DF is consistent
SearchResponse rsp = req.get();
assertSearchResponse(rsp);
SearchHits hits = rsp.getHits();
assertEquals(3, hits.getTotalHits());
assertEquals("1", hits.getAt(0).getId());
assertEquals("3", hits.getAt(1).getId());
assertEquals("2", hits.getAt(2).getId());
}
@Test
public void testScoring() throws Exception {
int numDocs = indexRandomNumbers("keyword");
String[] scoreModes = new String[] {"max", "min", "avg", "total", "multiply", ""};
float primaryWeight = 1.1f;
float secondaryWeight = 1.6f;
for (String scoreMode : scoreModes) {
for (int i = 0; i < numDocs - 4; i++) {
String[] intToEnglish =
new String[] {
English.intToEnglish(i),
English.intToEnglish(i + 1),
English.intToEnglish(i + 2),
English.intToEnglish(i + 3)
};
QueryRescorer rescoreQuery =
RescoreBuilder.queryRescorer(
QueryBuilders.boolQuery()
.disableCoord(true)
.should(
QueryBuilders.functionScoreQuery(
QueryBuilders.termQuery("field1", intToEnglish[0]))
.boostMode(CombineFunction.REPLACE)
.add(ScoreFunctionBuilders.scriptFunction("5.0f")))
.should(
QueryBuilders.functionScoreQuery(
QueryBuilders.termQuery("field1", intToEnglish[1]))
.boostMode(CombineFunction.REPLACE)
.add(ScoreFunctionBuilders.scriptFunction("7.0f")))
.should(
QueryBuilders.functionScoreQuery(
QueryBuilders.termQuery("field1", intToEnglish[3]))
.boostMode(CombineFunction.REPLACE)
.add(ScoreFunctionBuilders.scriptFunction("0.0f"))))
.setQueryWeight(primaryWeight)
.setRescoreQueryWeight(secondaryWeight);
if (!"".equals(scoreMode)) {
rescoreQuery.setScoreMode(scoreMode);
}
SearchResponse rescored =
client()
.prepareSearch()
.setPreference("test") // ensure we hit the same shards for tie-breaking
.setQuery(
QueryBuilders.boolQuery()
.disableCoord(true)
.should(
QueryBuilders.functionScoreQuery(
QueryBuilders.termQuery("field1", intToEnglish[0]))
.boostMode(CombineFunction.REPLACE)
.add(ScoreFunctionBuilders.scriptFunction("2.0f")))
.should(
QueryBuilders.functionScoreQuery(
QueryBuilders.termQuery("field1", intToEnglish[1]))
.boostMode(CombineFunction.REPLACE)
.add(ScoreFunctionBuilders.scriptFunction("3.0f")))
.should(
QueryBuilders.functionScoreQuery(
QueryBuilders.termQuery("field1", intToEnglish[2]))
.boostMode(CombineFunction.REPLACE)
.add(ScoreFunctionBuilders.scriptFunction("5.0f")))
.should(
QueryBuilders.functionScoreQuery(
QueryBuilders.termQuery("field1", intToEnglish[3]))
.boostMode(CombineFunction.REPLACE)
.add(ScoreFunctionBuilders.scriptFunction("0.2f"))))
.setFrom(0)
.setSize(10)
.setRescorer(rescoreQuery)
.setRescoreWindow(50)
.execute()
.actionGet();
assertHitCount(rescored, 4);
if ("total".equals(scoreMode) || "".equals(scoreMode)) {
assertFirstHit(rescored, hasId(String.valueOf(i + 1)));
assertSecondHit(rescored, hasId(String.valueOf(i)));
assertThirdHit(rescored, hasId(String.valueOf(i + 2)));
assertThat(
rescored.getHits().getHits()[0].getScore(),
equalTo(3.0f * primaryWeight + 7.0f * secondaryWeight));
assertThat(
rescored.getHits().getHits()[1].getScore(),
equalTo(2.0f * primaryWeight + 5.0f * secondaryWeight));
assertThat(rescored.getHits().getHits()[2].getScore(), equalTo(5.0f * primaryWeight));
assertThat(
rescored.getHits().getHits()[3].getScore(),
equalTo(0.2f * primaryWeight + 0.0f * secondaryWeight));
} else if ("max".equals(scoreMode)) {
assertFirstHit(rescored, hasId(String.valueOf(i + 1)));
assertSecondHit(rescored, hasId(String.valueOf(i)));
assertThirdHit(rescored, hasId(String.valueOf(i + 2)));
assertThat(rescored.getHits().getHits()[0].getScore(), equalTo(7.0f * secondaryWeight));
assertThat(rescored.getHits().getHits()[1].getScore(), equalTo(5.0f * secondaryWeight));
assertThat(rescored.getHits().getHits()[2].getScore(), equalTo(5.0f * primaryWeight));
assertThat(rescored.getHits().getHits()[3].getScore(), equalTo(0.2f * primaryWeight));
} else if ("min".equals(scoreMode)) {
assertFirstHit(rescored, hasId(String.valueOf(i + 2)));
assertSecondHit(rescored, hasId(String.valueOf(i + 1)));
assertThirdHit(rescored, hasId(String.valueOf(i)));
assertThat(rescored.getHits().getHits()[0].getScore(), equalTo(5.0f * primaryWeight));
assertThat(rescored.getHits().getHits()[1].getScore(), equalTo(3.0f * primaryWeight));
assertThat(rescored.getHits().getHits()[2].getScore(), equalTo(2.0f * primaryWeight));
assertThat(rescored.getHits().getHits()[3].getScore(), equalTo(0.0f * secondaryWeight));
} else if ("avg".equals(scoreMode)) {
assertFirstHit(rescored, hasId(String.valueOf(i + 1)));
assertSecondHit(rescored, hasId(String.valueOf(i + 2)));
assertThirdHit(rescored, hasId(String.valueOf(i)));
assertThat(
rescored.getHits().getHits()[0].getScore(),
equalTo((3.0f * primaryWeight + 7.0f * secondaryWeight) / 2.0f));
assertThat(rescored.getHits().getHits()[1].getScore(), equalTo(5.0f * primaryWeight));
assertThat(
rescored.getHits().getHits()[2].getScore(),
equalTo((2.0f * primaryWeight + 5.0f * secondaryWeight) / 2.0f));
assertThat(
rescored.getHits().getHits()[3].getScore(), equalTo((0.2f * primaryWeight) / 2.0f));
} else if ("multiply".equals(scoreMode)) {
assertFirstHit(rescored, hasId(String.valueOf(i + 1)));
assertSecondHit(rescored, hasId(String.valueOf(i)));
assertThirdHit(rescored, hasId(String.valueOf(i + 2)));
assertThat(
rescored.getHits().getHits()[0].getScore(),
equalTo(3.0f * primaryWeight * 7.0f * secondaryWeight));
assertThat(
rescored.getHits().getHits()[1].getScore(),
equalTo(2.0f * primaryWeight * 5.0f * secondaryWeight));
assertThat(rescored.getHits().getHits()[2].getScore(), equalTo(5.0f * primaryWeight));
assertThat(
rescored.getHits().getHits()[3].getScore(),
equalTo(0.2f * primaryWeight * 0.0f * secondaryWeight));
}
}
}
}
