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();
}
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));
}
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);
}
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));
}
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));
}
Aggregations