Search in sources :

Example 1 with AddDocumentResponse

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

the class ServerTestCase method addDocuments.

public static AddDocumentResponse addDocuments(Stream<AddDocumentRequest> requestStream) throws Exception {
    CountDownLatch finishLatch = new CountDownLatch(1);
    // observers responses from Server(should get one onNext and oneCompleted)
    final AtomicReference<AddDocumentResponse> response = new AtomicReference<>();
    final AtomicReference<Exception> exception = new AtomicReference<>();
    StreamObserver<AddDocumentResponse> responseStreamObserver = new StreamObserver<>() {

        @Override
        public void onNext(AddDocumentResponse value) {
            response.set(value);
        }

        @Override
        public void onError(Throwable t) {
            exception.set(new RuntimeException(t));
            finishLatch.countDown();
        }

        @Override
        public void onCompleted() {
            finishLatch.countDown();
        }
    };
    // requestObserver sends requests to Server (one onNext per AddDocumentRequest and one
    // onCompleted)
    StreamObserver<AddDocumentRequest> requestObserver = grpcServer.getStub().addDocuments(responseStreamObserver);
    // parse CSV into a stream of AddDocumentRequest
    try {
        requestStream.forEach(requestObserver::onNext);
    } catch (RuntimeException e) {
        // Cancel RPC
        requestObserver.onError(e);
        throw e;
    }
    // Mark the end of requests
    requestObserver.onCompleted();
    // Receiving happens asynchronously, so block here 20 seconds
    if (!finishLatch.await(20, TimeUnit.SECONDS)) {
        throw new RuntimeException("addDocuments can not finish within 20 seconds");
    }
    // Re-throw exception
    if (exception.get() != null) {
        throw exception.get();
    }
    return response.get();
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) AtomicReference(java.util.concurrent.atomic.AtomicReference) CountDownLatch(java.util.concurrent.CountDownLatch) AddDocumentRequest(com.yelp.nrtsearch.server.grpc.AddDocumentRequest) InvalidProtocolBufferException(com.google.protobuf.InvalidProtocolBufferException) IOException(java.io.IOException) AddDocumentResponse(com.yelp.nrtsearch.server.grpc.AddDocumentResponse)

Example 2 with AddDocumentResponse

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

the class ScoreScriptTest method testScriptUsingScoreInIndexField.

@Test
public void testScriptUsingScoreInIndexField() throws Exception {
    GrpcServer.TestServer testAddDocs = new GrpcServer.TestServer(grpcServer, false, Mode.STANDALONE);
    new GrpcServer.IndexAndRoleManager(grpcServer).createStartIndexAndRegisterFields(Mode.STANDALONE, 0, false, "registerFieldsScriptTest.json");
    AddDocumentResponse addDocumentResponse = testAddDocs.addDocuments("addDocs.csv");
    // manual refresh
    grpcServer.getBlockingStub().refresh(RefreshRequest.newBuilder().setIndexName(grpcServer.getTestIndex()).build());
    SearchResponse searchResponse = grpcServer.getBlockingStub().search(SearchRequest.newBuilder().setIndexName(grpcServer.getTestIndex()).addRetrieveFields("test_score").setStartHit(0).setTopHits(10).setQueryText("vendor_name:first vendor").build());
    assertEquals(2, searchResponse.getHitsCount());
    assertEquals(2.0, searchResponse.getHits(0).getFieldsOrThrow("test_score").getFieldValue(0).getDoubleValue(), Math.ulp(2.0));
    assertEquals(2.0, searchResponse.getHits(1).getFieldsOrThrow("test_score").getFieldValue(0).getDoubleValue(), Math.ulp(2.0));
}
Also used : GrpcServer(com.yelp.nrtsearch.server.grpc.GrpcServer) AddDocumentResponse(com.yelp.nrtsearch.server.grpc.AddDocumentResponse) SearchResponse(com.yelp.nrtsearch.server.grpc.SearchResponse) Test(org.junit.Test)

Example 3 with AddDocumentResponse

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

the class IndexerTask method index.

public Long index(LuceneServerClient luceneServerClient, Stream<AddDocumentRequest> addDocumentRequestStream) throws Exception {
    String threadId = Thread.currentThread().getName() + Thread.currentThread().getId();
    final CountDownLatch finishLatch = new CountDownLatch(1);
    StreamObserver<AddDocumentResponse> responseObserver = new StreamObserver<>() {

        @Override
        public void onNext(AddDocumentResponse value) {
            // Note that Server sends back only 1 message (Unary mode i.e. Server calls its onNext
            // only once
            // which is when it is done with indexing the entire stream), which means this method
            // should be
            // called only once.
            logger.debug("Received response for genId: {} on threadId: {}", value.getGenId(), threadId);
            genId = value.getGenId();
        }

        @Override
        public void onError(Throwable t) {
            logger.error(t.getMessage(), t);
            finishLatch.countDown();
        }

        @Override
        public void onCompleted() {
            logger.debug("Received final response from server on threadId: {}", threadId);
            finishLatch.countDown();
        }
    };
    // The responseObserver handles responses from the server (i.e. 1 onNext and 1 completed)
    // The requestObserver handles the sending of stream of client requests to server (i.e. multiple
    // onNext and 1 completed)
    StreamObserver<AddDocumentRequest> requestObserver = luceneServerClient.getAsyncStub().addDocuments(responseObserver);
    try {
        addDocumentRequestStream.forEach(addDocumentRequest -> requestObserver.onNext(addDocumentRequest));
    } catch (RuntimeException e) {
        // Cancel RPC
        requestObserver.onError(e);
        throw e;
    }
    // Mark the end of requests
    requestObserver.onCompleted();
    logger.debug("sent async addDocumentsRequest to server on threadId: {}", threadId);
    // Receiving happens asynchronously, so block here for 5 minutes
    if (!finishLatch.await(5, TimeUnit.MINUTES)) {
        logger.warn("addDocuments can not finish within 5 minutes on threadId: {}", threadId);
    }
    return Long.valueOf(genId);
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) CountDownLatch(java.util.concurrent.CountDownLatch) AddDocumentRequest(com.yelp.nrtsearch.server.grpc.AddDocumentRequest) AddDocumentResponse(com.yelp.nrtsearch.server.grpc.AddDocumentResponse)

Example 4 with AddDocumentResponse

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

the class ScoreScriptTest method testScriptDocValuesIndexField.

@Test
public void testScriptDocValuesIndexField() throws Exception {
    GrpcServer.TestServer testAddDocs = new GrpcServer.TestServer(grpcServer, false, Mode.STANDALONE);
    new GrpcServer.IndexAndRoleManager(grpcServer).createStartIndexAndRegisterFields(Mode.STANDALONE, 0, false, "registerFieldsScriptTest.json");
    AddDocumentResponse addDocumentResponse = testAddDocs.addDocuments("addDocs.csv");
    // manual refresh
    grpcServer.getBlockingStub().refresh(RefreshRequest.newBuilder().setIndexName(grpcServer.getTestIndex()).build());
    SearchResponse searchResponse = grpcServer.getBlockingStub().search(SearchRequest.newBuilder().setIndexName(grpcServer.getTestIndex()).addRetrieveFields("test_doc_values").setStartHit(0).setTopHits(10).build());
    assertEquals(2, searchResponse.getHitsCount());
    assertEquals(1.5, searchResponse.getHits(0).getFieldsOrThrow("test_doc_values").getFieldValue(0).getDoubleValue(), Math.ulp(1.5));
    assertEquals(1.5, searchResponse.getHits(1).getFieldsOrThrow("test_doc_values").getFieldValue(0).getDoubleValue(), Math.ulp(1.5));
}
Also used : GrpcServer(com.yelp.nrtsearch.server.grpc.GrpcServer) AddDocumentResponse(com.yelp.nrtsearch.server.grpc.AddDocumentResponse) SearchResponse(com.yelp.nrtsearch.server.grpc.SearchResponse) Test(org.junit.Test)

Example 5 with AddDocumentResponse

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

the class ScoreScriptTest method testQueryFieldScript.

private void testQueryFieldScript(String source, String registerFieldsFile, String addDocsFile, double expectedScore) throws Exception {
    GrpcServer.TestServer testAddDocs = new GrpcServer.TestServer(grpcServer, false, Mode.STANDALONE);
    new GrpcServer.IndexAndRoleManager(grpcServer).createStartIndexAndRegisterFields(Mode.STANDALONE, 0, false, registerFieldsFile);
    AddDocumentResponse addDocumentResponse = testAddDocs.addDocuments(addDocsFile);
    // manual refresh
    grpcServer.getBlockingStub().refresh(RefreshRequest.newBuilder().setIndexName(grpcServer.getTestIndex()).build());
    VirtualField virtualField = VirtualField.newBuilder().setName("test_field").setScript(Script.newBuilder().setLang("test_lang").setSource(source).build()).build();
    SearchResponse searchResponse = grpcServer.getBlockingStub().search(SearchRequest.newBuilder().setIndexName(grpcServer.getTestIndex()).setStartHit(0).setTopHits(10).addVirtualFields(virtualField).addRetrieveFields("test_field").build());
    assertEquals(2, searchResponse.getHitsCount());
    assertEquals(expectedScore, searchResponse.getHits(0).getFieldsOrThrow("test_field").getFieldValue(0).getDoubleValue(), Math.ulp(expectedScore));
    assertEquals(expectedScore, searchResponse.getHits(1).getFieldsOrThrow("test_field").getFieldValue(0).getDoubleValue(), Math.ulp(expectedScore));
}
Also used : VirtualField(com.yelp.nrtsearch.server.grpc.VirtualField) GrpcServer(com.yelp.nrtsearch.server.grpc.GrpcServer) AddDocumentResponse(com.yelp.nrtsearch.server.grpc.AddDocumentResponse) SearchResponse(com.yelp.nrtsearch.server.grpc.SearchResponse)

Aggregations

AddDocumentResponse (com.yelp.nrtsearch.server.grpc.AddDocumentResponse)5 GrpcServer (com.yelp.nrtsearch.server.grpc.GrpcServer)3 SearchResponse (com.yelp.nrtsearch.server.grpc.SearchResponse)3 AddDocumentRequest (com.yelp.nrtsearch.server.grpc.AddDocumentRequest)2 StreamObserver (io.grpc.stub.StreamObserver)2 CountDownLatch (java.util.concurrent.CountDownLatch)2 Test (org.junit.Test)2 InvalidProtocolBufferException (com.google.protobuf.InvalidProtocolBufferException)1 VirtualField (com.yelp.nrtsearch.server.grpc.VirtualField)1 IOException (java.io.IOException)1 AtomicReference (java.util.concurrent.atomic.AtomicReference)1