dhruva_RP-13176_support_multiple_indices_in_AddDocumentRequest (#862)

Added support for multiple indexNames in the same AddDocumentRequest. Created a new test file and updated NrtsearchClientBuilder and AddDocumentsCommand to support the new change.

Author: dhruva-r

clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto

@@ -832,7 +832,8 @@ message StartIndexResponse {
 
 // Message representing a document to index
 message AddDocumentRequest {
-    // Index name
+    // Index name; use this parameter when you want to add documents to one specific index.
+    // You must exclusively use this parameter or the indexNames parameter, not both.
     string indexName = 1;
     // Wrapper object to represent each field as a multivalued field.
     message MultiValuedField {
@@ -845,6 +846,9 @@ message AddDocumentRequest {
     map<string, MultiValuedField> fields = 3;
     // request type
     IndexingRequestType requestType = 4;
+    // List of indexNames; use this parameter when you want to add documents to multiple indices at once.
+    // You must exclusively use this parameter or the indexName parameter, not both.
+    repeated string indexNames = 5;
 }
 
 // Path for hierarchical facets

src/main/java/com/yelp/nrtsearch/server/grpc/NrtsearchClientBuilder.java

@@ -119,11 +119,15 @@ public StartIndexRequest buildRequest(Path filePath) throws IOException {
   class AddDocumentsClientBuilder implements NrtsearchClientBuilder<Stream<AddDocumentRequest>> {
     private static final Logger logger =
         LoggerFactory.getLogger(AddDocumentsClientBuilder.class.getName());
-    private final String indexName;
+    private final List<String> indexNames;
     private final CSVParser csvParser;
 
     public AddDocumentsClientBuilder(String indexName, CSVParser csvParser) {
-      this.indexName = indexName;
+      this(List.of(indexName), csvParser);
+    }
+
+    public AddDocumentsClientBuilder(List<String> indexNames, CSVParser csvParser) {
+      this.indexNames = indexNames;
       this.csvParser = csvParser;
     }
 
@@ -132,9 +136,9 @@ public Stream<AddDocumentRequest> buildRequest(Path filePath) {
       Stream.Builder<AddDocumentRequest> builder = Stream.builder();
       AddDocumentRequest.Builder addDocumentRequestBuilder = AddDocumentRequest.newBuilder();
       int cnt = 0;
+      addDocumentRequestBuilder.addAllIndexNames(indexNames);
       for (CSVRecord csvRecord : csvParser) {
         { // data rows
-          addDocumentRequestBuilder.setIndexName(indexName);
           for (String fieldName : csvParser.getHeaderNames()) {
             String fieldValues = csvRecord.get(fieldName);
             List<String> fieldVals = Arrays.asList(fieldValues.split(";"));
@@ -156,7 +160,7 @@ class AddJsonDocumentsClientBuilder
       implements NrtsearchClientBuilder<Stream<AddDocumentRequest>> {
     private static final Logger logger =
         LoggerFactory.getLogger(AddJsonDocumentsClientBuilder.class.getName());
-    private final String indexName;
+    private final List<String> indexNames;
     private final Gson gson;
     private final Path filePath;
     private final BufferedReader reader;
@@ -165,7 +169,12 @@ class AddJsonDocumentsClientBuilder
 
     public AddJsonDocumentsClientBuilder(
         String indexName, Gson gson, Path filePath, int maxBufferLen) throws IOException {
-      this.indexName = indexName;
+      this(List.of(indexName), gson, filePath, maxBufferLen);
+    }
+
+    public AddJsonDocumentsClientBuilder(
+        List<String> indexNames, Gson gson, Path filePath, int maxBufferLen) throws IOException {
+      this.indexNames = indexNames;
       this.gson = gson;
       this.filePath = filePath;
       this.reader = Files.newBufferedReader(filePath);
@@ -195,7 +204,7 @@ private void updateMultiValuedFielduilder(
     AddDocumentRequest buildAddDocumentRequest(String lineStr) {
       Map<String, Object> line = gson.fromJson(lineStr, Map.class);
       AddDocumentRequest.Builder addDocumentRequestBuilder = AddDocumentRequest.newBuilder();
-      addDocumentRequestBuilder.setIndexName(indexName);
+      addDocumentRequestBuilder.addAllIndexNames(indexNames);
       for (String key : line.keySet()) {
         AddDocumentRequest.MultiValuedField.Builder multiValuedFieldsBuilder =
             AddDocumentRequest.MultiValuedField.newBuilder();

src/main/java/com/yelp/nrtsearch/server/handler/AddDocumentHandler.java

@@ -97,21 +97,7 @@ private ArrayBlockingQueue<AddDocumentRequest> getAddDocumentRequestQueue(String
         }
       }
 
-      private long getCount(String indexName) {
-        return countMap.getOrDefault(indexName, 0L);
-      }
-
-      private void incrementCount(String indexName) {
-        if (countMap.containsKey(indexName)) {
-          countMap.put(indexName, countMap.get(indexName) + 1);
-        } else {
-          countMap.put(indexName, 1L);
-        }
-      }
-
-      @Override
-      public void onNext(AddDocumentRequest addDocumentRequest) {
-        String indexName = addDocumentRequest.getIndexName();
+      private void processIndexDocument(AddDocumentRequest request, String indexName) {
         ArrayBlockingQueue<AddDocumentRequest> addDocumentRequestQueue;
         try {
           addDocumentRequestQueue = getAddDocumentRequestQueue(indexName);
@@ -124,15 +110,14 @@ public void onNext(AddDocumentRequest addDocumentRequest) {
                 "onNext, index: %s, addDocumentRequestQueue size: %s",
                 indexName, addDocumentRequestQueue.size()));
         incrementCount(indexName);
-        addDocumentRequestQueue.add(addDocumentRequest);
+        addDocumentRequestQueue.add(request);
         if (addDocumentRequestQueue.remainingCapacity() == 0) {
           logger.debug(
               String.format(
                   "indexing addDocumentRequestQueue size: %s, total: %s",
                   addDocumentRequestQueue.size(), getCount(indexName)));
           try {
             DeadlineUtils.checkDeadline("addDocuments: onNext", "INDEXING");
-
             List<AddDocumentRequest> addDocRequestList = new ArrayList<>(addDocumentRequestQueue);
             Future<Long> future =
                 getGlobalState()
@@ -150,6 +135,44 @@ public void onNext(AddDocumentRequest addDocumentRequest) {
         }
       }
 
+      private long getCount(String indexName) {
+        return countMap.getOrDefault(indexName, 0L);
+      }
+
+      private void incrementCount(String indexName) {
+        if (countMap.containsKey(indexName)) {
+          countMap.put(indexName, countMap.get(indexName) + 1);
+        } else {
+          countMap.put(indexName, 1L);
+        }
+      }
+
+      @Override
+      public void onNext(AddDocumentRequest addDocumentRequest) {
+        String indexName = addDocumentRequest.getIndexName();
+        int indexNamesCount = addDocumentRequest.getIndexNamesCount();
+        if (indexName.isEmpty() && indexNamesCount == 0) {
+          onError(
+              Status.INVALID_ARGUMENT
+                  .withDescription(
+                      "Must provide exactly one of indexName or indexNames but neither is set")
+                  .asRuntimeException());
+        } else if (!indexName.isEmpty() && indexNamesCount > 0) {
+          onError(
+              Status.INVALID_ARGUMENT
+                  .withDescription(
+                      "Must provide exactly one of indexName or indexNames but both are set")
+                  .asRuntimeException());
+        } else if (indexNamesCount > 0) {
+          List<String> indexNames = addDocumentRequest.getIndexNamesList();
+          for (String currentIndexName : indexNames) {
+            processIndexDocument(addDocumentRequest, currentIndexName);
+          }
+        } else {
+          processIndexDocument(addDocumentRequest, indexName);
+        }
+      }
+
       @Override
       public void onError(Throwable t) {
         logger.warn("addDocuments Cancelled", t);

src/main/java/com/yelp/nrtsearch/tools/cli/AddDocumentsCommand.java

@@ -25,6 +25,7 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.List;
 import java.util.concurrent.Callable;
 import java.util.stream.Stream;
 import org.apache.commons.csv.CSVFormat;
@@ -60,12 +61,13 @@ public String getFileType() {
 
   @CommandLine.Option(
       names = {"-i", "--indexName"},
-      description = "Name of the index to add documents to",
+      description =
+          "List of index names to add documents to. Delimited by a \",\" : <indexName1>,<indexName2>,<indexName3>...",
       required = true)
-  private String indexName;
+  private String indexNamesStr;
 
-  public String getIndexName() {
-    return indexName;
+  public String getIndexNamesStr() {
+    return indexNamesStr;
   }
 
   @CommandLine.Option(
@@ -83,8 +85,8 @@ public int getMaxBufferLen() {
   public Integer call() throws Exception {
     NrtsearchClient client = baseCmd.getClient();
     try {
-      String indexName = getIndexName();
       String fileType = getFileType();
+      List<String> indexNames = List.of(getIndexNamesStr().split(","));
       Stream<AddDocumentRequest> addDocumentRequestStream;
       Path filePath = Paths.get(getFileName());
       if (fileType.equalsIgnoreCase("csv")) {
@@ -93,13 +95,13 @@ public Integer call() throws Exception {
             new CSVParser(
                 reader, CSVFormat.DEFAULT.builder().setHeader().setSkipHeaderRecord(true).build());
         addDocumentRequestStream =
-            new NrtsearchClientBuilder.AddDocumentsClientBuilder(indexName, csvParser)
+            new NrtsearchClientBuilder.AddDocumentsClientBuilder(indexNames, csvParser)
                 .buildRequest(filePath);
         client.addDocuments(addDocumentRequestStream);
       } else if (fileType.equalsIgnoreCase("json")) {
         NrtsearchClientBuilder.AddJsonDocumentsClientBuilder addJsonDocumentsClientBuilder =
             new NrtsearchClientBuilder.AddJsonDocumentsClientBuilder(
-                indexName, new Gson(), filePath, getMaxBufferLen());
+                indexNames, new Gson(), filePath, getMaxBufferLen());
         while (!addJsonDocumentsClientBuilder.isFinished()) {
           addDocumentRequestStream = addJsonDocumentsClientBuilder.buildRequest(filePath);
           client.addDocuments(addDocumentRequestStream);