Cleanup Rescorer Tests & add random indexing

src/test/java/org/elasticsearch/search/rescore/QueryRescorerTests.java

@@ -22,7 +22,7 @@ package org.elasticsearch.search.rescore;
 
 
 import org.apache.lucene.util.English;
-import org.apache.lucene.util.LuceneTestCase.Slow;
+import org.elasticsearch.action.index.IndexRequestBuilder;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.common.lucene.search.function.CombineFunction;
@@ -32,6 +32,7 @@ import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.index.query.MatchQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
+import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilders;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.SearchHits;
 import org.elasticsearch.search.rescore.RescoreBuilder.QueryRescorer;
@@ -39,7 +40,6 @@ import org.elasticsearch.test.ElasticsearchIntegrationTest;
 import org.junit.Test;
 
 import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
-import static org.elasticsearch.index.query.functionscore.ScoreFunctionBuilders.scriptFunction;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.*;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.notNullValue;
@@ -199,17 +199,16 @@ public class QueryRescorerTests extends ElasticsearchIntegrationTest {
                         "type1",
                         jsonBuilder().startObject().startObject("type1").startObject("properties").startObject("field1")
                                 .field("analyzer", "whitespace").field("type", "string").endObject().endObject().endObject().endObject())
-                .setSettings(ImmutableSettings.settingsBuilder()).execute().actionGet();
+                .setSettings(ImmutableSettings.settingsBuilder().put("index.number_of_shards", between(1, 5)).put("index.number_of_replicas", between(0, 1))).execute().actionGet();
         ensureGreen();
-        int numDocs = atLeast(100);
 
+        int numDocs = atLeast(100);
+        IndexRequestBuilder[] docs = new IndexRequestBuilder[numDocs];
         for (int i = 0; i < numDocs; i++) {
-            client().prepareIndex("test", "type1", String.valueOf(i)).setSource("field1", English.intToEnglish(i)).execute().actionGet();
+            docs[i] = client().prepareIndex("test", "type1", String.valueOf(i)).setSource("field1", English.intToEnglish(i));
         }
 
-        flush();
-        optimize(); // make sure we don't have a background merge running
-        refresh();
+        indexRandom(true, docs);
         ensureGreen();
         final int iters = atLeast(50);
         for (int i = 0; i < iters; i++) {
@@ -353,125 +352,14 @@ public class QueryRescorerTests extends ElasticsearchIntegrationTest {
                         "type1",
                         jsonBuilder().startObject().startObject("type1").startObject("properties").startObject("field1")
                                 .field("index", "not_analyzed").field("type", "string").endObject().endObject().endObject().endObject())
-                .setSettings(ImmutableSettings.settingsBuilder()).execute().actionGet();
-        ensureGreen();
-        int numDocs = 1000;
-
-        for (int i = 0; i < numDocs; i++) {
-            client().prepareIndex("test", "type1", String.valueOf(i)).setSource("field1", English.intToEnglish(i)).execute().actionGet();
-        }
-
-        flush();
-        optimize(); // make sure we don't have a background merge running
-        refresh();
-        ensureGreen();
-
-        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.customScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[0])).script("5.0f"))
-                                    .should(QueryBuilders.customScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[1])).script("7.0f"))
-                                    .should(QueryBuilders.customScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[3])).script("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.customScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[0])).script("2.0f"))
-                                .should(QueryBuilders.customScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[1])).script("3.0f"))
-                                .should(QueryBuilders.customScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[2])).script("5.0f"))
-                                .should(QueryBuilders.customScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[3])).script("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));
-                }
-            }
-        }
-    }
-
-
-
-    @Test
-    @Slow
-    public void testScoring_withFunctionScore() throws Exception {
-        client().admin()
-                .indices()
-                .prepareCreate("test")
-                .addMapping(
-                        "type1",
-                        jsonBuilder().startObject().startObject("type1").startObject("properties").startObject("field1")
-                                .field("index", "not_analyzed").field("type", "string").endObject().endObject().endObject().endObject())
-                .setSettings(ImmutableSettings.settingsBuilder()).execute().actionGet();
-        ensureGreen();
-        int numDocs = 1000;
-
+                .setSettings(ImmutableSettings.settingsBuilder().put("index.number_of_shards", between(1,5)).put("index.number_of_replicas", between(0,1))).get();
+        int numDocs = atLeast(100);
+        IndexRequestBuilder[] docs = new IndexRequestBuilder[numDocs];
         for (int i = 0; i < numDocs; i++) {
-            client().prepareIndex("test", "type1", String.valueOf(i)).setSource("field1", English.intToEnglish(i)).execute().actionGet();
+            docs[i] = client().prepareIndex("test", "type1", String.valueOf(i)).setSource("field1", English.intToEnglish(i));
         }
 
-        flush();
-        optimize(); // make sure we don't have a background merge running
-        refresh();
+        indexRandom(true, docs);
         ensureGreen();
 
         String[] scoreModes = new String[]{ "max", "min", "avg", "total", "multiply", "" };
@@ -485,10 +373,10 @@ public class QueryRescorerTests extends ElasticsearchIntegrationTest {
                 QueryRescorer rescoreQuery = RescoreBuilder
                         .queryRescorer(
                                 QueryBuilders.boolQuery()
-                                    .disableCoord(true)
-                                    .should(QueryBuilders.functionScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[0]), scriptFunction("5.0f")).boostMode(CombineFunction.REPLACE.getName()))
-                                    .should(QueryBuilders.functionScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[1]), scriptFunction("7.0f")).boostMode(CombineFunction.REPLACE.getName()))
-                                    .should(QueryBuilders.functionScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[3]), scriptFunction("0.0f")).boostMode(CombineFunction.REPLACE.getName())))
+                                        .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);
 
@@ -501,10 +389,10 @@ public class QueryRescorerTests extends ElasticsearchIntegrationTest {
                         .setPreference("test") // ensure we hit the same shards for tie-breaking
                         .setQuery(QueryBuilders.boolQuery()
                                 .disableCoord(true)
-                                .should(QueryBuilders.functionScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[0]), scriptFunction("2.0f")).boostMode(CombineFunction.REPLACE.getName()))
-                                .should(QueryBuilders.functionScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[1]), scriptFunction("3.0f")).boostMode(CombineFunction.REPLACE.getName()))
-                                .should(QueryBuilders.functionScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[2]), scriptFunction("5.0f")).boostMode(CombineFunction.REPLACE.getName()))
-                                .should(QueryBuilders.functionScoreQuery(QueryBuilders.termQuery("field1", intToEnglish[3]), scriptFunction("0.2f")).boostMode(CombineFunction.REPLACE.getName())))
+                                .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)
@@ -556,5 +444,4 @@ public class QueryRescorerTests extends ElasticsearchIntegrationTest {
             }
         }
     }
-
 }