Initialize Taxo Writer for Primary to support flat facets

src/main/java/com/yelp/nrtsearch/server/luceneserver/ShardState.java

@@ -478,38 +478,7 @@ public synchronized void start() throws Exception {
         openMode = IndexWriterConfig.OpenMode.APPEND;
       }
 
-      Path taxoDirFile;
-      if (rootDir == null) {
-        taxoDirFile = null;
-      } else {
-        taxoDirFile = rootDir.resolve("taxonomy");
-      }
-      taxoDir = indexState.df.open(taxoDirFile);
-
-      taxoSnapshots =
-          new PersistentSnapshotDeletionPolicy(
-              new KeepOnlyLastCommitDeletionPolicy(),
-              taxoDir,
-              IndexWriterConfig.OpenMode.CREATE_OR_APPEND);
-
-      taxoWriter =
-          new DirectoryTaxonomyWriter(taxoDir, openMode) {
-            @Override
-            protected IndexWriterConfig createIndexWriterConfig(
-                IndexWriterConfig.OpenMode openMode) {
-              IndexWriterConfig iwc = super.createIndexWriterConfig(openMode);
-              iwc.setIndexDeletionPolicy(taxoSnapshots);
-              return iwc;
-            }
-
-            @Override
-            protected IndexWriter openIndexWriter(Directory dir, IndexWriterConfig iwc)
-                throws IOException {
-              IndexWriter w = super.openIndexWriter(dir, iwc);
-              taxoInternalWriter = w;
-              return w;
-            }
-          };
+      initalizeTaxoWriter(openMode);
 
       writer =
           new IndexWriter(
@@ -625,6 +594,7 @@ public synchronized void startPrimary(long primaryGen, Path dataPath) throws Exc
       if (doCreate) {
         // nocommit shouldn't we set doCreate=false after we've done the create?
         openMode = IndexWriterConfig.OpenMode.CREATE;
+        initalizeTaxoWriter(openMode);
       } else {
         openMode = IndexWriterConfig.OpenMode.APPEND;
       }
@@ -704,6 +674,40 @@ public IndexSearcher newSearcher(IndexReader r, IndexReader previousReader)
     }
   }
 
+  private void initalizeTaxoWriter(IndexWriterConfig.OpenMode openMode) throws IOException {
+    Path taxoDirFile;
+    if (rootDir == null) {
+      taxoDirFile = null;
+    } else {
+      taxoDirFile = rootDir.resolve("taxonomy");
+    }
+    taxoDir = indexState.df.open(taxoDirFile);
+
+    taxoSnapshots =
+        new PersistentSnapshotDeletionPolicy(
+            new KeepOnlyLastCommitDeletionPolicy(),
+            taxoDir,
+            IndexWriterConfig.OpenMode.CREATE_OR_APPEND);
+
+    taxoWriter =
+        new DirectoryTaxonomyWriter(taxoDir, openMode) {
+          @Override
+          protected IndexWriterConfig createIndexWriterConfig(IndexWriterConfig.OpenMode openMode) {
+            IndexWriterConfig iwc = super.createIndexWriterConfig(openMode);
+            iwc.setIndexDeletionPolicy(taxoSnapshots);
+            return iwc;
+          }
+
+          @Override
+          protected IndexWriter openIndexWriter(Directory dir, IndexWriterConfig iwc)
+              throws IOException {
+            IndexWriter w = super.openIndexWriter(dir, iwc);
+            taxoInternalWriter = w;
+            return w;
+          }
+        };
+  }
+
   private IndexReader.ClosedListener removeSSDVStates =
       cacheKey -> {
         synchronized (ssdvStates) {

src/test/java/com/yelp/nrtsearch/server/luceneserver/ServerTestCase.java

@@ -64,6 +64,7 @@
  */
 public class ServerTestCase {
   public static final String DEFAULT_TEST_INDEX = "test_index";
+
   /**
    * This rule manages automatic graceful shutdown for the registered servers and channels at the
    * end of test.
@@ -226,6 +227,9 @@ private void initIndices() throws Exception {
       // start the index
       StartIndexRequest.Builder startIndexBuilder =
           StartIndexRequest.newBuilder().setIndexName(indexName);
+      if (shouldInitializeIndicesAsPrimary()) {
+        startIndexBuilder.setMode(Mode.PRIMARY);
+      }
       blockingStub.startIndex(startIndexBuilder.build());
 
       // add Docs
@@ -236,6 +240,10 @@ private void initIndices() throws Exception {
     }
   }
 
+  protected boolean shouldInitializeIndicesAsPrimary() {
+    return false;
+  }
+
   protected List<String> getIndices() {
     return Collections.singletonList(DEFAULT_TEST_INDEX);
   }

src/test/java/com/yelp/nrtsearch/server/luceneserver/facet/NumberFieldFlatFacetsTest.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2020 Yelp Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.yelp.nrtsearch.server.luceneserver.facet;
+
+import static org.junit.Assert.assertNotNull;
+
+import com.yelp.nrtsearch.server.grpc.AddDocumentRequest;
+import com.yelp.nrtsearch.server.grpc.FieldDefRequest;
+import com.yelp.nrtsearch.server.luceneserver.ServerTestCase;
+import io.grpc.testing.GrpcCleanupRule;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.junit.ClassRule;
+import org.junit.Test;
+
+public class NumberFieldFlatFacetsTest extends ServerTestCase {
+  @ClassRule public static final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+
+  private static final List<String> fields =
+      Arrays.asList(
+          new String[] {
+            "int_number_facet_field",
+            "float_number_facet_field",
+            "long_number_facet_field",
+            "double_number_facet_field"
+          });
+  private static final List<String> numericValues =
+      Arrays.asList(new String[] {"1", "10", "20", "30"});
+
+  private Map<String, AddDocumentRequest.MultiValuedField> getFieldsMapForOneDocument(
+      String value) {
+    Map<String, AddDocumentRequest.MultiValuedField> fieldsMap = new HashMap<>();
+    for (String field : fields) {
+      fieldsMap.put(
+          field, AddDocumentRequest.MultiValuedField.newBuilder().addValue(value).build());
+    }
+    return fieldsMap;
+  }
+
+  @Override
+  public FieldDefRequest getIndexDef(String name) throws IOException {
+    return getFieldsFromResourceFile("/facet/number_field_flat_facets.json");
+  }
+
+  @Override
+  protected boolean shouldInitializeIndicesAsPrimary() {
+    return true;
+  }
+
+  @Test
+  public void addDocuments() throws InterruptedException {
+    List<AddDocumentRequest> documentRequests = new ArrayList<>();
+    for (String value : numericValues) {
+      documentRequests.add(
+          AddDocumentRequest.newBuilder()
+              .setIndexName(DEFAULT_TEST_INDEX)
+              .putAllFields(getFieldsMapForOneDocument(value))
+              .build());
+    }
+    assertNotNull(addDocuments(documentRequests.stream()));
+  }
+}

src/test/resources/facet/number_field_flat_facets.json

@@ -0,0 +1,37 @@
+{
+  "indexName": "test_index",
+  "field": [
+    {
+      "name": "int_number_facet_field",
+      "type": "INT",
+      "storeDocValues": true,
+      "multiValued": false,
+      "search": true,
+      "facet": "FLAT"
+    },
+    {
+      "name": "float_number_facet_field",
+      "type": "FLOAT",
+      "storeDocValues": true,
+      "multiValued": false,
+      "search": true,
+      "facet": "FLAT"
+    },
+    {
+      "name": "long_number_facet_field",
+      "type": "LONG",
+      "storeDocValues": true,
+      "multiValued": false,
+      "search": true,
+      "facet": "FLAT"
+    },
+    {
+      "name": "double_number_facet_field",
+      "type": "DOUBLE",
+      "storeDocValues": true,
+      "multiValued": false,
+      "search": true,
+      "facet": "FLAT"
+    }
+  ]
+}