Examples with SearchRequest com.yelp.nrtsearch.server.grpc.SearchRequest used on opensource projects

Example 21 with SearchRequest

use of com.yelp.nrtsearch.server.grpc.SearchRequest in project nrtsearch by Yelp.

the class DocCollectorTest method testUsesDefaultTerminateAfter.

@Test
public void testUsesDefaultTerminateAfter() {
    IndexState indexState = Mockito.mock(IndexState.class);
    when(indexState.getDefaultTerminateAfter()).thenReturn(100);
    SearchRequest request = SearchRequest.newBuilder().setTopHits(10).build();
    TestDocCollector docCollector = new TestDocCollector(request, indexState);
    assertTrue(docCollector.getManager() instanceof TestDocCollector.TestCollectorManager);
    assertTrue(docCollector.getWrappedManager() instanceof TerminateAfterWrapper);
    assertEquals(100, ((TerminateAfterWrapper<?>) docCollector.getWrappedManager()).getTerminateAfter());
}

Example 22 with SearchRequest

use of com.yelp.nrtsearch.server.grpc.SearchRequest in project nrtsearch by Yelp.

the class SearchRequestProcessor method buildDocCollector.

/**
 * Build {@link DocCollector} to provide the {@link org.apache.lucene.search.CollectorManager} for
 * collecting hits for this query.
 *
 * @return collector
 */
private static DocCollector buildDocCollector(CollectorCreatorContext collectorCreatorContext) {
    SearchRequest searchRequest = collectorCreatorContext.getRequest();
    List<AdditionalCollectorManager<? extends Collector, ? extends CollectorResult>> additionalCollectors = searchRequest.getCollectorsMap().entrySet().stream().map(e -> CollectorCreator.getInstance().createCollectorManager(collectorCreatorContext, e.getKey(), e.getValue())).collect(Collectors.toList());
    if (searchRequest.getQuerySort().getFields().getSortedFieldsList().isEmpty()) {
        if (hasLargeNumHits(searchRequest)) {
            return new LargeNumHitsCollector(collectorCreatorContext, additionalCollectors);
        } else {
            return new RelevanceCollector(collectorCreatorContext, additionalCollectors);
        }
    } else {
        return new SortFieldCollector(collectorCreatorContext, additionalCollectors);
    }
}

Example 23 with SearchRequest

use of com.yelp.nrtsearch.server.grpc.SearchRequest in project nrtsearch by Yelp.

the class CollectorTest method testCustomCollector.

/**
 * Collect the top 3 reviews for each business *
 */
@Test
public void testCustomCollector() throws InvalidProtocolBufferException {
    // topHits = 0 results in java.lang.IllegalArgumentException:
    // numHits must be > 0; please use TotalHitCountCollector if you just need the total hit count
    SearchRequest request = SearchRequest.newBuilder().setTopHits(1).setIndexName(DEFAULT_TEST_INDEX).addRetrieveFields("review_id").setQuery(Query.newBuilder().setFunctionScoreQuery(FunctionScoreQuery.newBuilder().setScript(Script.newBuilder().setLang("js").setSource("review_count").build()).build()).build()).putCollectors(COLLECTOR_NAME, Collector.newBuilder().setPluginCollector(PluginCollector.newBuilder().setName(PLUGIN_COLLECTOR).setParams(Struct.newBuilder().putFields("size", Value.newBuilder().setNumberValue(3).build())).build()).build()).build();
    Map<String, CollectorResult> collectorsResult = getGrpcServer().getBlockingStub().search(request).getCollectorResultsMap();
    assertEquals(1, collectorsResult.size());
    // review_id -> score
    Map<String, Double> expectedHits = Map.of("0", 100.0, "1", 99.0, "2", 98.0, "3", 97.0, "4", 96.0, "5", 95.0, "6", 94.0, "7", 93.0, "8", 92.0);
    Map<String, Double> collectedHits = collectorsResult.get(COLLECTOR_NAME).getAnyResult().unpack(SearchResponse.class).getHitsList().stream().collect(Collectors.toMap(hit -> hit.getFieldsMap().get("review_id").getFieldValue(0).getTextValue(), SearchResponse.Hit::getScore));
    assertEquals(expectedHits, collectedHits);
}

Example 24 with SearchRequest

use of com.yelp.nrtsearch.server.grpc.SearchRequest in project nrtsearch by Yelp.

the class CollectorTest method testUnknownCollector.

@Test
public void testUnknownCollector() {
    SearchRequest request = SearchRequest.newBuilder().setTopHits(1).setIndexName(DEFAULT_TEST_INDEX).addRetrieveFields("review_id").setQuery(Query.newBuilder().setFunctionScoreQuery(FunctionScoreQuery.newBuilder().setScript(Script.newBuilder().setLang("js").setSource("review_count").build()).build()).build()).putCollectors(COLLECTOR_NAME, com.yelp.nrtsearch.server.grpc.Collector.newBuilder().setPluginCollector(PluginCollector.newBuilder().setName("invalid").build()).build()).build();
    try {
        getGrpcServer().getBlockingStub().search(request);
        fail();
    } catch (StatusRuntimeException e) {
        String expectedMessage = "Invalid collector name: invalid, must be one of: [plugin_collector]";
        assertTrue(e.getMessage().contains(expectedMessage));
    }
}

Example 25 with SearchRequest

use of com.yelp.nrtsearch.server.grpc.SearchRequest in project nrtsearch by Yelp.

the class DocCollectorTest method testUsesDefaultTimeout.

@Test
public void testUsesDefaultTimeout() {
    IndexState indexState = Mockito.mock(IndexState.class);
    when(indexState.getDefaultSearchTimeoutSec()).thenReturn(3.0);
    when(indexState.getDefaultSearchTimeoutCheckEvery()).thenReturn(0);
    SearchRequest request = SearchRequest.newBuilder().setTopHits(10).build();
    TestDocCollector docCollector = new TestDocCollector(request, indexState);
    assertTrue(docCollector.getManager() instanceof TestDocCollector.TestCollectorManager);
    assertTrue(docCollector.getWrappedManager() instanceof SearchCutoffWrapper);
    SearchCutoffWrapper<?> cutoffWrapper = (SearchCutoffWrapper<?>) docCollector.getWrappedManager();
    assertEquals(3.0, cutoffWrapper.getTimeoutSec(), 0.0);
    assertEquals(0, cutoffWrapper.getCheckEvery());
}